How can you use PyTorch Profiler on SageMaker?

Question

Is there a way to use the PyTorch profiler when using Amazon SageMaker? Or can we call PyTorch profiler in the profiler config for SageMaker debugger

Answer 1

You can use the below snippet to initiate a Profiler .

from sagemaker.pytorch import PyTorch
from sagemaker.debugger import ProfilerConfig, FrameworkProfile
profiler_config = ProfilerConfig(
    framework_profile_params=FrameworkProfile(start_step=1, num_steps=2)
)

Then you can pass the profilfer_config to the Pytorch estimator.

Answer 2

To use the PyTorch vanilla profiler in SageMaker Training jobs:

1. On the SageMaker Estimator turn off the SageMaker built-in profiler (as you'll be using the PT Vanilla profiler): estimator = PyTorch(..., disable_profiler=True, ...)
2. In your training script add PT profiler as shown in this PT Profiler with TensorBoard tutorial.
3. Export the profiling data to SageMaker managed output folder: on_trace_ready=torch.profiler.tensorboard_trace_handler(f'{args.output_data_dir}/profile')
4. when the job completes. Start TensorBoard with the S3 location of the output data folder: estimator.latest_job_tensorboard_artifacts_path()