SageMaker Profiler を使用したトレーニングジョブの準備と実行

SageMaker Profiler でトレーニングジョブを実行するためのセットアップは、トレーニングスクリプトの適応と SageMaker トレーニングジョブランチャーの設定の 2 つのステップで構成されます。

ステップ 1: SageMaker Profiler Python モジュールを使用したトレーニングスクリプトの調整

トレーニングジョブの実行中に GPU で実行されるカーネルのキャプチャを開始するには、SageMaker Profiler Python モジュールを使用してトレーニングスクリプトを変更します。ライブラリをインポートし、start_profiling() と stop_profiling() メソッドを追加して、プロファイリングの開始と終了を定義します。オプションのカスタム注釈を使用してトレーニングスクリプトにマーカーを追加し、各ステップの特定のオペレーション中のハードウェアアクティビティを視覚化することもできます。

アノテーターは GPU からオペレーションを抽出することに注意してください。CPU でのプロファイリングオペレーションでは、注釈を追加する必要はありません。CPU プロファイリングは、ステップ 2: SageMaker AI フレームワーク推定器の作成と SageMaker Profiler の有効化 で練習するプロファイリング設定を指定するときにも有効になります。

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 Profiler の有効化

以下の手順は、SageMaker Python SDK を使用してトレーニング用の SageMaker AI フレームワーク推定器を準備する方法を示しています。

トレーニングジョブの実行中またはジョブの完了後に、SageMaker Profiler 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 パッケージをインストールする

SageMaker Profiler がプリインストールされた SageMaker AI フレームワークイメージ にリストされていない PyTorch または TensorFlow フレームワークイメージ、またはトレーニング用の独自のカスタム Docker コンテナで SageMaker Profiler を使用するには、SageMaker Profiler Python パッケージバイナリファイル のいずれかを使用して SageMaker Profiler をインストールします。

オプション 1: トレーニングジョブの起動中に SageMaker Profiler パッケージをインストールする

SageMaker Profiler がプリインストールされた SageMaker AI フレームワークイメージ にリストされていない PyTorch または TensorFlow イメージを使用してジョブをトレーニングするために SageMaker Profiler を使用する場合は、requirements.txt ファイルを作成し、ステップ 2 で SageMaker AI フレームワーク推定器の source_dir パラメータに指定したパスの下に配置します。requirements.txt ファイル全般のセットアップの詳細については、「SageMaker Python SDK ドキュメント」の「Using third-party libraries」を参照してください。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 コンテナを使用する場合は、Dockerfile に SageMaker Profiler Python パッケージバイナリファイル のいずれかを追加します。

# 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

SageMaker AI でトレーニング用にカスタム Docker コンテナを実行する一般的なガイダンスについては、「独自のトレーニングコンテナの適応」を参照してください。