Amazon SageMaker、独自学習ループAPI

LLMのファインチューニングでは、通常、学習手順の制御と運用の簡便さのどちらかを選ばされてきた。フルマネージドサービスはベースモデル、データ、ハイパーパラメータ以外のすべてを管理するため、標準的なレシピが合う場合には有効だ。管理範囲の小さい選択肢はアルゴリズムの自由度が増すが、インフラ作業がチームに移る。Amazon SageMaker AIの学習セッションAPIは、アルゴリズムではなくコンピュートを抽象化することで、このトレードオフを解消しようとしている。

プログラミングモデルは、順に作成する3つのSDKクライアントを中心に構成される。プラットフォームレベルの設定用のServiceClient、学習とチェックポイント操作のTrainingClient、ロールアウト中の生成とスコア計算のSamplingClientだ。forward_backward()、optim_step()、sample()といった操作は、稼働中のステートフルなセッションに対するSageMaker API呼び出しとなり、ステップをまたいでモデルはGPU上に読み込まれたままになる。ランタイムのセットアップ、分散実行、勾配同期、障害復旧、テナント分離、大規模なGPU運用はサービスが担う。

クックブックは、最小ループを超えるエンジニアリングをパッケージ化している。サンプリングと学習のオーバーラップ、ストリーミング勾配計算、ホールドアウト検証、チェックポイント再開、ステップごとのメトリクスなどだ。code_rlレシピはマルチターンのツール利用を追加し、モデルは採点前にcheck_solutionツールを呼び出して回答を修正できる。生成されたPythonはSandboxFusionコンテナ内でユニットテストに対して実行される。投稿によれば、50ステップで学習報酬は0.46から0.60に上昇し、200問のホールドアウトテスト問題の精度は40パーセントから50.5パーセントに上昇した。