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

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

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

[torch.profiler](https://pytorch.org/docs/stable/profiler.html#torch.profiler.profile) と [PyTorch Profiler TensorBoard Plugin](https://github.com/pytorch/kineto/tree/main/tb_plugin) は、それぞれ公式に詳しく利用方法が記載されていますので、あわせてご確認ください。

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

Profiling data を生成する学習コードは、 [amazon-sagemaker-examples/sagemaker-python-sdk/pytorch_mnist/](https://github.com/aws/amazon-sagemaker-examples/tree/main/sagemaker-python-sdk/pytorch_mnist) を利用しました。 PyTorch Profiler TensorBoard Plugin のリポジトリに[サンプルコード](https://github.com/pytorch/kineto/tree/main/tb_plugin/examples)が用意されていますので、こちらも参考にしています。

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](https://docs.aws.amazon.com/sagemaker/latest/dg/model-train-storage.html)
* trace の step 数が多くなると、後述の SageMaker Studio から TensorBoard を起動する際に読み込みに失敗しますので、 `schedule=torch.profiler.schedule(wait, warmup, active, repeat, skip_first)` の `repeat` で調整しましょう。

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

```diff_python
--- 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 を実行します。

```python
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](https://github.com/pytorch/pytorch/issues/69443) が発生したため、 `1.12.0` を利用しています。 

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

```python
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]` を起動して、コマンドを実行します。
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/9216/133a360d-2301-f408-ee19-2572f7398f28.png)

先程の 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://docs.aws.amazon.com/sagemaker/latest/dg/studio-tensorboard.html)

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

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

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

### Overview
![image (1).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/9216/f215c337-096d-0c4f-b94a-9dbd8397dac3.png)

### Operator
![image (2).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/9216/17c3764f-fe6e-a6ea-3f72-77a2caba1929.png)

### GPU Kernel
![image (3).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/9216/47703d5f-580b-30e9-49d8-9e6acc57185c.png)

### Trace
![image (4).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/9216/5c0ab465-5c51-b10a-0806-aa4d50d091fe.png)

### Memory
![image (5).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/9216/e13d0693-3d9e-4c78-b41f-32f539e3e988.png)