Optimun Neuron を使うと、HuggingFace TrainerをAWS Neuron上で気軽に扱えるようになります。具体的には以下のように、TrainerとTrainingArgumentsを置き換えるだけです。
from optimum.neuron import NeuronTrainer as Trainer from optimum.neuron import NeuronTrainingArguments as TrainingArguments
こうするとTrn1上で、今までと同じコードで学習が可能です。AWS Neuronを使う場合はモデルのコンパイルが必要なのですが、これも自動で行われます。学習の間に動的にコンパイルを実行し、学習をしてくれます。この動的なコンパイルは長くて30分ほどかかりますが、学習自体が速いため総合的な時間は短縮されるでしょう。
作成されるモデルはコンパイル済みのものではありません。AWS Neuronというとそれ専用になるイメージがありますが、大丈夫です。
スクリプトの実行時は torchrun というコマンドを使います。
torchrun --nproc_per_node=2 train_model.py
公式のサンプルノートブックがおすすめ
こちらに公式で公開されているノートブックが用意されています。やりたいことに合わせて参考にするのがおすすめです！
ただNotebookであれこれ試すとなるとNeuronのコンパイルが頻繁に走って面倒な場合があります。個人的には手元のGPUなどでしっかりチューニングしてから、AWS Neuronで学習させるのが良いと思います。データセットの準備なども手元で済ませるほうが良いです（これについては別途記事に）。
SageMakerでは勝手にtorchrunをしてくれる
SageMakerでTrn1を使う場合は、自動で train_model.py を torchrun で実行してくれます。ですのでやることといえばインポートを切り替えるくらいなのでとても便利です。
モデルのコンパイル時間にどれくらいかかるか？などもモデルによりけりなので、ぜひ一度試してみると良いでしょう。インスタンスタイプを変えるだけで、勝手にSageMakerがよしなにやってくれます。
他にも知っておくと良いかも、ということ
AWS Neuronでの学習について詳しく知りたいとき
Optimun Neuronを使うと良い感じにラップされすぎていて、逆に何をしているのか分かりにくいときがあります。
pytorh-xlaを使ってNeuron上でTrainerを使った学習のサンプルがAWSの公式であるのでこちらもおすすめです。
この記事、以前のインタビューでもご一緒した常世さんが書いてらっしゃいます。
trn1.32xlargeでマルチノードを使う場合は設定が必要
Trn1でマルチノードで学習させる際は、EFAの設定が必要になります。SecurityGroupで所属するマシン同士ですべてのInbound、Outboundを許可したり、専用のネットワークインタフェースを作成する必要があります。EC2の場合、手動で準備する必要がありかなり手間は大きいです。以下のドキュメントでコマンドによる作成方法などが書かれていますのでこちらを参考にしてください。
awsdocs-neuron.readthedocs-hosted.com
設定がうまくいっているかチェックするコマンドがとても役立ちます。ネットワークのPing Pongも確認してくれるので、SecurityGroupの設定が正しいかなども確認できます（実際にこのコマンドのおかげで間違いに気づいたこともありました）。
チェックポイントから再開させるとバグ
Optimun Neuronのバージョンによってはチェックポイントから再開する際にバグが発生する可能性があります。細かいことは次回の記事で記載いたします。
おわりに
AWS NeuronはTrnもInfも情報が少ないですが、公式できちんと整備されたドキュメントやサンプルはたくさんあります。ぜひお手元のモデルでの学習に活用してみてください。ShodoではAWS Neuronも最大限活用し、AIを使ったサービス運営を行っています。
現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80％オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をブログの執筆にお役立てください。
XMAS2024
