Amazon SageMaker Training と Amazon SageMaker Studio で PyTorch Profiler TensorBoard Plugin を利用する方法

機械学習においても Observability 👀 大切ですよね！

PyTorch を利用した機械学習において、ボトルネック特定のために PyTorch Profiler TensorBoard Plugin を活用する機会がありましたので、SageMaker Training と SageMaker Studio から利用する方法をまとめます。

torch.profiler と PyTorch Profiler TensorBoard Plugin は、それぞれ公式に詳しく利用方法が記載されていますので、あわせてご確認ください。

SageMaker Training のコードで Profiling data を生成する

Profiling data を生成する学習コードは、 amazon-sagemaker-examples/sagemaker-python-sdk/pytorch_mnist/ を利用しました。 PyTorch Profiler TensorBoard Plugin のリポジトリにサンプルコードが用意されていますので、こちらも参考にしています。

SageMaker Training で Profiling data を生成する際の注意点は以下の 2 つです。

• 環境変数 SM_OUTPUT_DATA_DIR(=/opt/ml/output/data/) 以下の任意のディレクトリを出力先 on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name) に指定します。参考: Amazon SageMaker Training Storage Folders for Training Datasets, Checkpoints, Model Artifacts, and Outputs
• trace の step 数が多くなると、後述の SageMaker Studio から TensorBoard を起動する際に読み込みに失敗しますので、 schedule=torch.profiler.schedule(wait, warmup, active, repeat, skip_first) の repeat で調整しましょう。

以下は mnist.py 変更箇所の diff です。

--- a/sagemaker-python-sdk/pytorch_mnist/mnist.py
+++ b/sagemaker-python-sdk/pytorch_mnist/mnist.py
@@ -10,6 +10,7 @@ import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
+import torch.profiler
 import torch.utils.data
 import torch.utils.data.distributed
 from torchvision import datasets, transforms
@@ -140,31 +141,53 @@ def train(args):
 
     optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
 
+    tensorboard_dir = f'{args.output_data_dir}/tensorboard'
+    logger.info(f'tensorboard_dir: {tensorboard_dir}')
+    
     for epoch in range(1, args.epochs + 1):
         model.train()
-        for batch_idx, (data, target) in enumerate(train_loader, 1):
-            data, target = data.to(device), target.to(device)
-            optimizer.zero_grad()
-            output = model(data)
-            loss = F.nll_loss(output, target)
-            loss.backward()
-            if is_distributed and not use_cuda:
-                # average gradients manually for multi-machine cpu case only
-                _average_gradients(model)
-            optimizer.step()
-            if batch_idx % args.log_interval == 0:
-                logger.info(
-                    "Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}".format(
-                        epoch,
-                        batch_idx * len(data),
-                        len(train_loader.sampler),
-                        100.0 * batch_idx / len(train_loader),
-                        loss.item(),
+        
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA],
+            schedule=torch.profiler.schedule(
+                wait=0,
+                warmup=1,
+                active=5,
+                repeat=10),
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(tensorboard_dir),
+            record_shapes=True,
+            profile_memory=True,
+            with_stack=False,
+        ) as prof:
+            
+            for batch_idx, (data, target) in enumerate(train_loader, 1):
+                data, target = data.to(device), target.to(device)
+                optimizer.zero_grad()
+                output = model(data)
+                loss = F.nll_loss(output, target)
+                loss.backward()
+                if is_distributed and not use_cuda:
+                    # average gradients manually for multi-machine cpu case only
+                    _average_gradients(model)
+                optimizer.step()
+                prof.step()
+                if batch_idx % args.log_interval == 0:
+                    logger.info(
+                        "Train Epoch: {} [{}/{} ({:.0f}%)] Loss: {:.6f}".format(
+                            epoch,
+                            batch_idx * len(data),
+                            len(train_loader.sampler),
+                            100.0 * batch_idx / len(train_loader),
+                            loss.item(),
+                        )
                     )
-                )
         test(model, test_loader, device)
     save_model(model, args.model_dir)
 
+    tensorboard_files = os.listdir(tensorboard_dir)
+    logger.info(f'tensorboard_files: {tensorboard_files}')
 
 def test(model, test_loader, device):
     model.eval()
@@ -253,5 +276,6 @@ if __name__ == "__main__":
     parser.add_argument("--model-dir", type=str, default=os.environ["SM_MODEL_DIR"])
     parser.add_argument("--data-dir", type=str, default=os.environ["SM_CHANNEL_TRAINING"])
     parser.add_argument("--num-gpus", type=int, default=os.environ["SM_NUM_GPUS"])
+    parser.add_argument("--output-data-dir", type=str, default=os.environ['SM_OUTPUT_DATA_DIR'])
 
     train(parser.parse_args())

では、SageMaker Training を実行します。

from sagemaker.pytorch import PyTorch

estimator = PyTorch(
    entry_point="mnist.py",
    role=role,
    py_version="py38",
    framework_version="1.12.0",
    instance_count=1,
    instance_type="ml.p3.2xlarge",
    hyperparameters={"epochs": 1, "backend": "gloo"},
)

estimator.fit({"training": inputs})

参考: 1.11 の PyTorch Profiler は SEGFAULT が発生したため、 1.12.0 を利用しています。

Profiling data が出力されたことがログからも確認できました。

INFO:__main__:tensorboard_files: ['algo-1_43.1670765907589.pt.trace.json', 'algo-1_43.1670765907378.pt.trace.json', 'algo-1_43.1670765908323.pt.trace.json', 'algo-1_43.1670765907887.pt.trace.json', 'algo-1_43.1670765907155.pt.trace.json', 'algo-1_43.1670765906381.pt.trace.json', 'algo-1_43.1670765908612.pt.trace.json', 'algo-1_43.1670765906918.pt.trace.json', 'algo-1_43.1670765906599.pt.trace.json', 'algo-1_43.1670765908112.pt.trace.json']

SageMaker Studio から TensorBoard を起動する

SageMaker Studio の左上アイコンをクリックすると表示される [Launcher] の [Utilities and files] より [System terminal] を起動して、コマンドを実行します。

先程の SageMaker Python SDK のコードで PyTorch Estimator の output_path は未指定のため、出力先は SageMaker のデフォルト S3 バケットです。output.tar.gz が保存されていることを SageMaker Studio の System terminal から aws-cli で確認します。

sagemaker-user@studio$ aws s3 ls s3://sagemaker-[region]-[account id]/[Training job name]-[time]/output/
2022-12-11 13:38:56      82334 model.tar.gz
2022-12-11 13:38:57     709989 output.tar.gz

大丈夫ですね。

では、PyTorch Profiler TensorBoard Plugin をインストールします。

sagemaker-user@studio$ pip install torch-tb-profiler
...
Successfully installed torch-tb-profiler-0.4.0

続いて、s3 から output.tar.gz を取得して、 TensorBoard を起動します。

sagemaker-user@studio$ aws s3 cp s3://sagemaker-[region]-[account id]/[Training job name]-[time]/output/output.tar.gz /tmp/
...

sagemaker-user@studio$ rm -rf /tmp/mnist/ && mkdir -p /tmp/mnist/ && tar zxvf /tmp/output.tar.gz -C /tmp/mnist/
tensorboard/
tensorboard/algo-1_43.1670765907589.pt.trace.json
tensorboard/algo-1_43.1670765907378.pt.trace.json
tensorboard/algo-1_43.1670765908323.pt.trace.json
tensorboard/algo-1_43.1670765907887.pt.trace.json
tensorboard/algo-1_43.1670765907155.pt.trace.json
tensorboard/algo-1_43.1670765906381.pt.trace.json
tensorboard/algo-1_43.1670765908612.pt.trace.json
tensorboard/algo-1_43.1670765906918.pt.trace.json
tensorboard/algo-1_43.1670765906599.pt.trace.json
tensorboard/algo-1_43.1670765908112.pt.trace.json

sagemaker-user@studio$ tensorboard --logdir=/tmp/mnist/tensorboard
...
TensorBoard 2.11.0 at http://localhost:6006/ (Press CTRL+C to quit)
...

SageMaker Studio から TensorBoard を起動するには、以下のように Studio の URL をコピーして、lab 以下を proxy/6006/ に置き換えます。参考: Use TensorBoard in Amazon SageMaker Studio

https://[your url].studio.[region].sagemaker.aws/jupyter/default/proxy/6006/

PyTorch Profiler TensorBoard Plugin のスクリーンショット

最後にスクリーンショットを貼り付けておきます。どなたかのご参考になることを祈っておりますー！

Overview

Operator

GPU Kernel

Trace

Memory