Find bottlenecks in your code (advanced)¶
Audience: Users who want to profile their TPU models to find bottlenecks and improve performance.
Profile cloud TPU models¶
To profile TPU models use the XLAProfiler
from lightning.pytorch.profilers import XLAProfiler
profiler = XLAProfiler(port=9001)
trainer = Trainer(profiler=profiler)
Capture profiling logs in Tensorboard¶
To capture profile logs in Tensorboard, follow these instructions:
0: Setup the required installs¶
Use this guide to help you with the Cloud TPU required installations.
1: Start Tensorboard¶
Start the TensorBoard server:
tensorboard --logdir ./tensorboard --port 9001
Now open the following url on your browser
http://localhost:9001/#profile
2: Capture the profile¶
Once the code you want to profile is running:
click on the
CAPTURE PROFILE
button.Enter
localhost:9001
(default port for XLA Profiler) as the Profile Service URL.Enter the number of milliseconds for the profiling duration
Click
CAPTURE
3: Don’t stop your code¶
Make sure the code is running while you are trying to capture the traces. It will lead to better performance insights if the profiling duration is longer than the step time.
4: View the profiling logs¶
Once the capture is finished, the page will refresh and you can browse through the insights using the Tools dropdown at the top left