Amazon SageMaker AI、訓練ループを自作できるtraining session APIsを追加
Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs
AWSはAmazon SageMaker AIのサーバーレスモデルカスタマイズにtraining session APIsを追加した。forward_backward()、optim_step()、sample()をAPI呼び出しとして組み合わせ、SFTや強化学習の訓練ループを手元のPythonで書ける。GPUの割り当てや分散実行、障害復旧はSageMaker AIが担う。
training session APIsは、アルゴリズムではなく計算資源を抽象化する点が特徴である。ライブでステートフルなセッションにモデルがGPU上に載ったまま、ユーザーが必要な回数だけ呼び出しを重ねる。ベースモデルとLoRAアダプターが読み込まれ、実行時のセットアップ、分散実行、勾配同期、障害復旧、テナント分離はサービス側が処理する。クライアントはServiceClient、TrainingClient、SamplingClientの3種類で、アクティブなセッションはOpenAI互換のチャット補完エンドポイントも提供する。
課金はトークン処理量に基づく。報酬計算やデータ読み込み、デバッグが手元で行われている間はGPU時間が課金されず、アイドル時間のコストが生じない。複雑な報酬モデルやマルチツールのエージェント型ロールアウトなど、ローカル計算の重い処理で差が出やすいとされる。位置付けとしては、設定のみで行うServerless Model Customizationより制御度が高く、SageMaker Training JobsやHyperPodより運用負担が小さい。
記事ではGSM8Kを使い、OpenAIのGPT-OSS 20BのLoRAアダプターをGRPOで訓練する最小ループを解説している。サンプリング、ローカルでの報酬採点、勾配更新、サンプラー更新の4段階で構成され、各GPU呼び出しはfutureを返すため並列に処理できる。チェックポイントはモデルパッケージとして登録される。このサンプルでは、10ステップ、40問でホールドアウト精度が56%から75%に上がった。
実運用に近い例として、cookbookのcode_rlレシピを紹介している。DeepCoderデータセットの競技プログラミング問題を使い、生成コードをSandboxFusionコンテナで単体テストにかけて採点する。2ターンでcheck_solutionツールを呼び出せ、サンプリングと訓練を重ねる非同期実行も備える。GPT-OSS 20Bを50ステップ訓練した結果、訓練報酬は0.46から0.60、200問のテストでの精度は40%から50.5%になった。cookbookにはほかにTulu 3やUltraChatでのSFT、KTOによる選好アラインメントなどのレシピもある。
出典: AWS Machine Learning Blog · aws.amazon.com