Use Amazon SageMaker Profiler to profile activities on Amazon compute resources

Amazon SageMaker Profiler is a capability of Amazon SageMaker that provides a detailed view into the Amazon compute resources provisioned during training deep learning models on SageMaker. It focuses on profiling the CPU and GPU usage, kernel runs on GPUs, kernel launches on CPUs, sync operations, memory operations across CPUs and GPUs, latencies between kernel launches and corresponding runs, and data transfer between CPUs and GPUs. SageMaker Profiler also offers a user interface (UI) that visualizes the profile, a statistical summary of profiled events, and the timeline of a training job for tracking and understanding the time relationship of the events between GPUs and CPUs.

For data scientists

Training deep learning models on a large compute cluster often has computational optimization problems, such as bottlenecks, kernel launch latencies, memory limit, and low resource utilization.

To identify such computational performance issues, you need to profile deeper into the compute resources to understand which kernels introduce latencies and which operations cause bottlenecks. Data scientists can take the benefit from using the SageMaker Profiler UI for visualizing the detailed profile of training jobs. The UI provides a dashboard furnished with summary charts and a timeline interface to track every event on the compute resources. Data scientists can also add custom annotations to track certain parts of the training job using the SageMaker Profiler Python modules.

For administrators

Through the Profiler landing page in the SageMaker console or SageMaker domain, you can manage the Profiler application users if you are an administrator of an Amazon account or SageMaker domain. Each domain user can access their own Profiler application given the granted permissions. As a SageMaker domain administrator and domain user, you can create and delete the Profiler application given the permission level you have.

Supported framework images, Amazon Web Services Regions, and instance types

This feature supports the following machine learning frameworks and Amazon Web Services Regions.

SageMaker framework images pre-installed with SageMaker Profiler

SageMaker Profiler is pre-installed in the following Amazon Deep Learning Containers for SageMaker.

PyTorch images

TensorFlow images

SageMaker Profiler Python package binary files

If you want to configure your own Docker container, use SageMaker Profiler in other pre-built containers for PyTorch and TensorFlow, or install the SageMaker Profiler Python package locally, use one the following binary files. Depending on the Python and CUDA versions in your environment, choose one of the following.

PyTorch

TensorFlow

For more information about how to install SageMaker Profiler using the binary files, see (Optional) Install the SageMaker Profiler Python package.

Supported Amazon Web Services Regions

SageMaker Profiler is available in the following Amazon Web Services Regions.

Supported instance types

SageMaker Profiler supports profiling of training jobs on the following instance types.

CPU and GPU profiling

GPU profiling only

Prerequisites

The following list shows the prerequisites to start using SageMaker Profiler.

Prepare and run a training job with SageMaker Profiler

Setting up to running a training job with the SageMaker Profiler consists of two steps: adapting the training script and configuring the SageMaker training job launcher.

Step 1: Adapt your training script using the SageMaker Profiler Python modules

To start capturing kernel runs on GPUs while the training job is running, modify your training script using the SageMaker Profiler Python modules. Import the library and add the start_profiling() and stop_profiling() methods to define the beginning and the end of profiling. You can also use optional custom annotations to add markers in the training script to visualize hardware activities during particular operations in each step.

Note that the annotators extract operations from GPUs. For profiling operations in CPUs, you don’t need to add any additional annotations. CPU profiling is also activated when you specify the profiling configuration, which you’ll practice in Step 2: Create a SageMaker framework estimator and activate SageMaker Profiler.

try:
    import smprof 
except ImportError:
    # backward-compatability for TF 2.11 and PT 1.13.1 images
    import smppy as smprof

Approach 1. Use the context manager smprof.annotate to annotate full functions

You can wrap full functions with the smprof.annotate() context manager. This wrapper is recommended if you want to profile by functions instead of code lines. The following example script shows how to implement the context manager to wrap the training loop and full functions in each iteration.

import smprof

SMProf = smprof.SMProfiler.instance()
config = smprof.Config()
config.profiler = {
    "EnableCuda": "1",
}
SMProf.configure(config)
SMProf.start_profiling()

for epoch in range(args.epochs):
    if world_size > 1:
        sampler.set_epoch(epoch)
    tstart = time.perf_counter()
    for i, data in enumerate(trainloader, 0):
        with smprof.annotate("step_"+str(i)):
            inputs, labels = data
            inputs = inputs.to("cuda", non_blocking=True)
            labels = labels.to("cuda", non_blocking=True)
    
            optimizer.zero_grad()
    
            with smprof.annotate("Forward"):
                outputs = net(inputs)
            with smprof.annotate("Loss"):
                loss = criterion(outputs, labels)
            with smprof.annotate("Backward"):
                loss.backward()
            with smprof.annotate("Optimizer"):
                optimizer.step()

SMProf.stop_profiling()

Approach 2. Use smprof.annotation_begin() and smprof.annotation_end() to annotate specific code line in functions

You can also define annotations to profile specific code lines. You can set the exact starting point and end point of profiling at the level of individual code lines, not by the functions. For example, in the following script, the step_annotator is defined at the beginning of each iteration and ends at the end of the iteration. Meanwhile, other detailed annotators for each operations are defined and wrap around the target operations throughout each iteration.

import smprof

SMProf = smprof.SMProfiler.instance()
config = smprof.Config()
config.profiler = {
    "EnableCuda": "1",
}
SMProf.configure(config)
SMProf.start_profiling()

for epoch in range(args.epochs):
    if world_size > 1:
        sampler.set_epoch(epoch)
    tstart = time.perf_counter()
    for i, data in enumerate(trainloader, 0):
        step_annotator = smprof.annotation_begin("step_" + str(i))

        inputs, labels = data
        inputs = inputs.to("cuda", non_blocking=True)
        labels = labels.to("cuda", non_blocking=True)
        optimizer.zero_grad()

        forward_annotator = smprof.annotation_begin("Forward")
        outputs = net(inputs)
        smprof.annotation_end(forward_annotator)

        loss_annotator = smprof.annotation_begin("Loss")
        loss = criterion(outputs, labels)
        smprof.annotation_end(loss_annotator)

        backward_annotator = smprof.annotation_begin("Backward")
        loss.backward()
        smprof.annotation_end(backward_annotator)

        optimizer_annotator = smprof.annotation_begin("Optimizer")
        optimizer.step()
        smprof.annotation_end(optimizer_annotator)

        smprof.annotation_end(step_annotator)

SMProf.stop_profiling()

After annotating and setting up the profiler initiation modules, save the script to submit using a SageMaker training job launcher in the following Step 2. The sample launcher assumes that the training script is named train_with_profiler_demo.py.

Step 2: Create a SageMaker framework estimator and activate SageMaker Profiler

The following procedure shows how to prepare a SageMaker framework estimator for training using the SageMaker Python SDK.

While running the training job or after the job has completed, you can go to the next topic at Open the SageMaker Profiler UI application and start exploring and visualizing the saved profiles.

If you want to directly access the profile data saved in the Amazon S3 bucket, use the following script to retrieve the S3 URI.

import os
# This is an ad-hoc function to get the S3 URI
# to where the profile output data is saved
def get_detailed_profiler_output_uri(estimator):
    config_name = None
    for processing in estimator.profiler_rule_configs:
        params = processing.get("RuleParameters", dict())
        rule = config_name = params.get("rule_to_invoke", "")
        if rule == "DetailedProfilerProcessing":
            config_name = processing.get("RuleConfigurationName")
            break
    return os.path.join(
        estimator.output_path, 
        estimator.latest_training_job.name, 
        "rule-output",
        config_name,
    )

print(
    f"Profiler output S3 bucket: ", 
    get_detailed_profiler_output_uri(estimator)
)

(Optional) Install the SageMaker Profiler Python package

To use SageMaker Profiler on PyTorch or TensorFlow framework images not listed in SageMaker framework images pre-installed with SageMaker Profiler, or on your own custom Docker container for training, you can install SageMaker Profiler by using one of the SageMaker Profiler Python package binary files.

Option 1: Install the SageMaker Profiler package while launching a training job

If you want to use SageMaker Profiler for training jobs using PyTorch or TensorFlow images not listed in SageMaker framework images pre-installed with SageMaker Profiler, create a requirements.txt file and locate it under the path you specify to the source_dir parameter of the SageMaker framework estimator in Step 2. For more information about setting up a requirements.txt file in general, see Using third-party libraries in the SageMaker Python SDK documentation. In the requirements.txt file, add one of the S3 bucket paths for the SageMaker Profiler Python package binary files.

# requirements.txt
https://smppy.s3.amazonaws.com/tensorflow/cu112/smprof-0.3.332-cp39-cp39-linux_x86_64.whl

Option 2: Install the SageMaker Profiler package in your custom Docker containers

If you use a custom Docker container for training, add one of the SageMaker Profiler Python package binary files to your Dockerfile.

# Install the smprof package version compatible with your CUDA version
RUN pip install https://smppy.s3.amazonaws.com/tensorflow/cu112/smprof-0.3.332-cp39-cp39-linux_x86_64.whl

For guidance on running a custom Docker container for training on SageMaker in general, see Adapting your own training container.

Open the SageMaker Profiler UI application

You can access the SageMaker Profiler UI application through the following options.

Option 1: Launch the SageMaker Profiler UI from the domain details page

If you have access to the SageMaker console, you can take this option.

Navigate to the domain details page

The following procedure shows how to navigate to the domain details page.

Launch the SageMaker Profiler UI application

The following procedure shows how to launch the SageMaker Profiler application that is scoped to a user profile.

Option 2: Launch the SageMaker Profiler UI application from the SageMaker Profiler landing page in the SageMaker console

The following procedure describes how to launch the SageMaker Profiler UI application from the SageMaker Profiler landing page in the SageMaker console. If you have access to the SageMaker console, you can take this option.

Option 3: Use the application launcher function in the SageMaker Python SDK

If you are a SageMaker domain user and have access only to SageMaker Studio, you can access the SageMaker Profiler UI application through SageMaker Studio Classic by running the sagemaker.interactive_apps.detail_profiler_app.DetailProfilerApp function.

Note that SageMaker Studio Classic is the previous Studio UI experience before re:Invent 2023, and is migrated as an application into a newly designed Studio UI at re:Invent 2023. The SageMaker Profiler UI application is available at SageMaker domain level, and thus requires your domain ID and user profile name. Currently, the DetailedProfilerApp function only works within the SageMaker Studio Classic application; the function properly takes in the domain and user profile information from SageMaker Studio Classic.

For domain, domain users, and Studio created before re:Invent 2023, Studio Classic would be the default experience unless you have updated it following the instructions at Migrating from Amazon SageMaker Studio Classic. If this is your case, there's no further action needed, and you can directly launch the SageMaker Profiler UI application by running the DetailProfilerApp funciton.

If you created a new domain and Studio after re:Invent 2023, launch the Studio Classic application within the Studio UI and then run the DetailProfilerApp function to launch the SageMaker Profiler UI application.

Note that the DetailedProfilerApp function doesn’t work in other SageMaker machine learning IDEs, such as the SageMaker Studio JupyterLab application, the SageMaker Studio Code Editor application, and SageMaker Notebook instances. If you run the DetailedProfilerApp function in those IDEs, it returns a URL to the Profiler landing page in the SageMaker console, instead of a direct link to open the Profiler UI application.

Explore the profile output data visualized in the SageMaker Profiler UI

This section walks through the SageMaker Profiler UI and provides tips for how to use and gain insights from it.

Load profile

When you open the SageMaker Profiler UI, the Load profile page opens up. To load and generate the Dashboard and Timeline, go through the following procedure.

Dashboard

After you finish loading and selecting the training job, the UI opens the Dashboard page furnished with the following panels by default.

The following screenshots show the GPU and CPU active time ratio and the average GPU and CPU utilization rate with respect to time per compute node.

The following screenshot shows an example of pie charts for comparing how many times the GPU kernels are launched and measuring the time spent on running them. In the Time spent by all GPU kernels and Launch counts of all GPU kernels panels, you can also specify an integer to the input field for k to adjust the number of legend to show in the plots. For example, if you specify 10, the plots show the top 10 most run and launched kernels respectively.

The following screenshot shows an example of step time duration histogram, and pie charts for the kernel precision distribution, GPU activity distribution, and GPU memory operation distribution.

Timeline interface

To gain a detailed view into the compute resources at the level of operations and kernels scheduled on the CPUs and run on the GPUs, use the Timeline interface.

You can zoom in and out and pan left or right in the timeline interface using your mouse, the [w, a, s, d] keys, or the four arrow keys on the keyboard.

The timeline tracks are organized in a tree structure, giving you information from the host level to the device level. For example, if you run N instances with eight GPUs in each, the timeline structure of each instance would be as follows.

The following several screenshots show the Timeline of the profile of a training job run on ml.p4d.24xlarge instances, which are equipped with 8 NVIDIA A100 Tensor Core GPUs in each.

The following is a zoomed-out view of the profile, printing a dozen of steps including an intermittent data loader between step_232 and step_233 for fetching the next data batch.

For each CPU, you can track the CPU utilization and performance counters, such as "clk_unhalted_ref.tsc" and "itlb_misses.miss_causes_a_walk", which are indicative of instructions run on the CPU.

For each GPU, you can see a host timeline and a device timeline. Kernel launches are on the host timeline and kernel runs are on the device timeline. You can also see annotations (such as forward, backward, and optimize) if you have added in training script in the GPU host timeline.

In the timeline view, you can also track kernel launch-and-run pairs. This helps you understand how a kernel launch scheduled on a host (CPU) is run on the corresponding GPU device.

The following screenshot is a zoomed-in view into step_233 and step_234 from the previous screenshot. The timeline interval selected in the following screenshot is the AllReduce operation, an essential communication and synchronization step in distributed training, run on the GPU-0 device. In the screenshot, note that the kernel launch in the GPU-0 host connects to the kernel run in the GPU-0 device stream 1, indicated with the arrow in cyan color.

Also two information tabs appear in the bottom pane of the UI when you select a timeline interval, as shown in the previous screenshot. The Current Selection tab shows the details of the selected kernel and the connected kernel launch from the host. The connection direction is always from host (CPU) to device (GPU) since each GPU kernel is always called from a CPU. The Connections tab shows the chosen kernel launch and run pair. You can select either of them to move it to the center of the Timeline view.

The following screenshot zooms in further into the AllReduce operation launch and run pair.

Information

In Information, you can access information about the loaded training job, such as the instance type, Amazon Resource Names (ARNs) of compute resources provisioned for the job, node names, and hyperparameters.

Settings

The SageMaker Profiler UI application instance is configured to shut down after 2 hours of idle time by default. In Settings, use the following settings to adjust the auto shutdown timer.

Frequently asked questions about using SageMaker Profiler

Use the following frequently asked questions to find answers about using SageMaker Profiler.

Q. I’m getting an error message, ModuleNotFoundError: No module named 'smppy'

Since December 2023, the name of the SageMaker Profiler Python package has changed from smppy to smprof to resolve a duplicate package name issue; smppy is already used by an open source package.

Therefore, if you have been using smppy since before December 2023 and experiencing this ModuleNotFoundError issue, it might be due to the outdated package name in your training script while having the latested smprof package installed or using one of the latest SageMaker framework images pre-installed with SageMaker Profiler. In this case, make sure that you replace all mentions of smppy with smprof throughout your training script.

While updating the SageMaker Profiler Python package name in your training scripts, to avoid confusion around which version of the package name you should use, consider using a conditional import statement as shown in the following code snippet.

try:
    import smprof 
except ImportError:
    # backward-compatability for TF 2.11 and PT 1.13.1 images
    import smppy as smprof

Also note that if you have been using smppy while upgrading to the latest PyTorch or TensorFlow versions, make sure that you install the latest smprof package by following instructions at (Optional) Install the SageMaker Profiler Python package.

Q. I’m getting an error message, ModuleNotFoundError: No module named 'smprof'

First, make sure that you use one of the officially supported SageMaker Framework Containers. If you don’t use one of those, you can install the smprof package by following instructions at (Optional) Install the SageMaker Profiler Python package.

Q. I’m not able to import ProfilerConfig

If you are unable to import ProfilerConfig in your job launcher script using the SageMaker Python SDK, your local environment or the Jupyter kernel might have a significantly outdated version of the SageMaker Python SDK. Make sure that you upgrade the SDK to the latest version.

$ pip install --upgrade sagemaker

Q. I’m getting an error message, aborted: core dumped when importing smprof into my training script

In an earlier version of smprof, this issue occurs with PyTorch 2.0+ and PyTorch Lightning. To resolve this issue, also install the latest smprof package by following instructions at (Optional) Install the SageMaker Profiler Python package.

Q. I cannot find the SageMaker Profiler UI from SageMaker Studio. How can I find it?

If you have access to the SageMaker console, choose one of the following options.

If you are a domain user and don't have access to the SageMaker console, you can access the application through SageMaker Studio Classic. If this is your case, choose the following option.

Considerations

Consider the following when using SageMaker Profiler.