使用 SageMaker 探查器准备和运行训练作业

使用 SageMaker 探查器设置以运行训练作业包括两个步骤：调整训练脚本和配置 SageMaker 训练作业启动器。

步骤 1：使用 SageMaker 探查器 Python 模块调整训练脚本

要在训练作业运行时开始捕获 GPU 上内核运行，请使用 SageMaker 探查器 Python 模块修改训练脚本。导入库并添加 start_profiling() 和 stop_profiling() 方法来定义分析的开始和结束。您还可以使用可选的自定义注释在训练脚本中添加标记，以便可视化每个步骤中的特定操作期间的硬件活动。

请注意，注释器会从 GPU 中提取操作。对于 CPU 中的分析操作，您无需添加任何其他注释。在指定分析配置（您将在步骤 2：创建 SageMaker AI 框架估算器并激活 SageMaker 探查器中使用此配置）时也将激活 CPU 分析。

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

方法 1. 使用上下文管理器 smprof.annotate 为所有函数添加注释

您可以通过 smprof.annotate() 上下文管理器封装所有函数。如果您想按函数而不是代码行进行分析，建议使用此包装器。以下示例脚本说明如何实施上下文管理器，以便在每次迭代中包装训练循环和完整函数。

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()

方法 2. 使用 smprof.annotation_begin() 和 smprof.annotation_end() 为函数中的特定代码行添加注释

您还可以定义注释来分析特定的代码行。可以在单个代码行级别而不是按函数来设置分析的确切起始点和结束点。例如，在以下脚本中，step_annotator 会在每次迭代开始时被定义，并在迭代结束时终止。同时，还为每个操作定义了其他详细注释器，并在每次迭代中包装目标操作。

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()

在注释并设置探查器初始模块后，在接下来的步骤 2 中，使用 SageMaker 训练作业启动器保存脚本以进行提交。示例启动器假定训练脚本已命名为 train_with_profiler_demo.py。

步骤 2：创建 SageMaker AI 框架估算器并激活 SageMaker 探查器

以下过程说明如何使用 SageMaker Python SDK 准备 SageMaker AI 框架估算器以进行训练。

在运行训练作业时或作业完成后，您可以转到 打开 SageMaker 探查器 UI 应用程序 上的下一个主题，并开始探究和可视化已保存的配置文件。

如果您想直接访问保存在 Amazon S3 存储桶中的配置文件数据，请使用以下脚本来检索 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)
)

（可选）安装 SageMaker Profiler Python 软件包

要在 PyTorch 或 TensorFlow 框架上使用 SageMaker Profiler，或者在自己的自定义 Docker 容器上进行训练，而 预安装了 SageMaker 探查器的 SageMaker AI 框架映像 中未列出这些映像，您可以使用 SageMaker Profiler Python 软件包二进制文件 之一来安装 SageMaker Profiler。

选项 1：在启动训练作业时安装 SageMaker Profiler 软件包

如果您想使用 SageMaker 探查器处理训练作业，但这些作业使用的是未在预安装了 SageMaker 探查器的 SageMaker AI 框架映像中列出的 PyTorch 或 TensorFlow 映像，请创建一个 requirements.txt 文件并将其放置在步骤 2 中 SageMaker AI 框架估计器的 source_dir 参数指定的路径下。有关设置 requirements.txt 文件的更多信息，请参阅 SageMaker Python SDK 文档中的使用第三方库。在 requirements.txt 文件中，为 SageMaker Profiler Python 软件包二进制文件 文件添加一个 S3 存储桶路径。

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

选项 2：在自定义 Docker 容器中安装 SageMaker Profiler 软件包

如果您使用自定义 Docker 容器进行训练，请将 SageMaker Profiler Python 软件包二进制文件 之一添加到您的 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

有关运行自定义 Docker 容器以进行 SageMaker AI 训练的一般指导，请参阅 Adapting your own training container。