AWS Trainiumを使用した事前学習済みモデルのコンパイル
AWS Inferentiaを活用するモデルを使用するには、AWS Neuron SDKを使用してAWS Inferentia用にコンパイルする必要があります。
これはInferentia用にモデルをコンパイルするために使用するコードです：
import torch
import numpy as np
import os
import torch_neuronx
from torchvision import models
image = torch.zeros([1, 3, 224, 224], dtype=torch.float32)
## Load a pretrained ResNet50 model
model = models.resnet50(pretrained=True)
## Tell the model we are using it for evaluation (not training)
model.eval()
model_neuron = torch_neuronx.trace(model, image)
## Export to saved model
model_neuron.save("resnet50_neuron.pt")
このコードは事前学習済みのResNet-50モデルを読み込み、評価モードに設定します。ここではモデルにトレーニングデータを追加していないことに注意してください。その後、AWS Neuron SDKを使用してモデルを保存します。
EKSクラスタにPodをデプロイし、AWS Inferentia用のサンプルモデルをコンパイルします。AWS Inferentia用のモデルをコンパイルするには、AWS Neuron SDKが必要です。このSDKはAWSが提供するDeep Learning Containers (DLCs)に含まれています。
デバイスプラグインのインストール
DLCがNeuronコアを使用するためには、それらを公開する必要があります。NeuronデバイスプラグインのkubernetesマニフェストファイルはNeuronコアをDLCに公開します。これらのマニフェストファイルはEKSクラスタに事前にインストールされています。
Podが公開されたNeuronコアを必要とする場合、Kubernetesスケジューラは、Podをスケジュールするために、InferentiaまたはTrainiumノードをプロビジョニングできます。
実行する予定のイメージを確認します：
トレーニング用のPodを作成する
このコードをEKS上のPodで実行します。これはPodを実行するためのマニフェストファイルです：
apiVersion: v1
kind: Pod
metadata:
labels:
role: compiler
name: compiler
namespace: aiml
spec:
nodeSelector:
node.kubernetes.io/instance-type: trn1.2xlarge
containers:
- command:
- sh
- -c
- sleep infinity
image: ${AIML_DL_TRN_IMAGE}
name: compiler
resources:
limits:
aws.amazon.com/neuron: 1
serviceAccountName: inference
nodeSelectorセクションでは、このPodを実行したいインスタンスタイプを指定しています。この場合はtrn1インスタンスです。
resourcesのlimitsセクションでは、このPodを実行するためにneuronコアが必要であることを指定しています。これによりNeuronデバイスプラグインはPodにneuron APIを公開するように指示します。
次のコマンドを実行してPodを作成します：
Karpenterはtrn1インスタンスとNeuronコアを必要とする保留中のPodを検出し、要件を満たすtrn1インスタンスを起動します。次のコマンドを使用してインスタンスのプロビジョニングを監視できます：
{"level": "INFO",
"time": "2024-09-19T18:44:08.919Z",
"logger": "controller",
"message": "launched nodeclaim",
"commit": "6e9d95f",
"controller": "nodeclaim.lifecycle",
"controllerGroup": "karpenter.sh",
"controllerKind": "NodeClaim",
"NodeClaim": {"name": "aiml-hp9wm"
},
"namespace": "",
"name": "aiml-hp9wm",
"reconcileID": "b38f0b3c-f146-4544-8ddc-ca73574c97f0",
"provider-id": "aws:///us-west-2b/i-06bc9a7cb6f92887c",
"instance-type": "trn1.2xlarge",
"zone": "us-west-2b",
"capacity-type": "on-demand",
"allocatable": {"aws.amazon.com/neuron": "1",
"cpu": "7910m",
"ephemeral-storage": "89Gi",
"memory": "29317Mi",
"pods": "58",
"vpc.amazonaws.com/pod-eni": "17"
}
}
PodはKarpenterによってプロビジョニングされたノードにスケジュールされるはずです。Podが準備完了の状態になっているか確認します：
このコマンドは最大で10分かかる場合があります。
次に、モデルをコンパイルするためのコードをPodにコピーして実行します：
....
Downloading: "https://download.pytorch.org/models/resnet50-0676ba61.pth" to /root/.cache/torch/hub/checkpoints/resnet50-0676ba61.pth
100%|-------| 97.8M/97.8M [00:00<00:00, 165MB/s]
.
Compiler status PASS
最後に、モデルを作成済みのS3バケットにアップロードします。これにより、後でラボでモデルを使用できるようになります。
upload: ./resnet50_neuron.pt to s3://eksworkshop-inference20230511204343601500000001/resnet50_neuron.pt