今日、大規模言語モデル(LLM)のファインチューニングでは、通常、学習ループの制御と運用の手軽さの間でトレードオフが生じます。フルマネージド型のサービスでは、ベースモデルを選択し、データとハイパーパラメータを提供すれば、残りはすべてサービス側が担ってくれます。これは、標準的な既製のレシピが問題に合っている場合にはうまく機能します。管理度の低いサービスではより多くの制御が得られますが、その分、基盤となるインフラストラクチャや学習フレームワークの管理もより多く担うことになります。標準的なレシピでは不十分な場合もあります。教師ありファインチューニング(SFT)用のカスタム損失やデータカリキュラム、独自の選好最適化、強化学習(RL)の報酬関数、あるいは先週発表されたアライメントアルゴリズムが必要になることがあります。学習手順を自ら形作る必要が生じた瞬間に、アルゴリズムの自由を得るためにそのインフラ作業を引き受けることになります。
本日、Amazon SageMaker AI のサーバーレスモデルカスタマイズが、training session API によって拡張されました。この API は抽象化の境界を一段下のレベルに引き下げることで、そのトレードオフを解消します。アルゴリズムを抽象化するのではなく、コンピューティングを抽象化します。例えば次のような操作です。 forward_backward(), optim_step(), および sample() お客様の手法に必要な手順——SFT、選好最適化、RL、あるいは全く新しい手法——を組み合わせて構成するSageMaker API呼び出しになります。これらの呼び出しはライブでステートフルなセッションを対象とするため、モデルはステップをまたいでGPU上にロードされたまま維持され、必要な数の呼び出しを自由に実行できます。ランタイムのセットアップ、分散実行、勾配同期、障害復旧、テナント分離、そして大規模なGPU運用はサービス側が処理します。お客様はトレーニングループを記述し、インフラはSageMaker AIが実行します。
このサービスとあわせて、私たちは以下を公開します。 SageMakerトレーニングセッションクックブック、すぐに実行できるSFTおよびRLレシピのセットであり、独自のトレーニングループの設計図として使用できます。
この記事では、プログラミングモデルと3つのSDKクライアントを紹介し、SageMaker AIのモデルトレーニングおよびカスタマイズのオプションの中でトレーニングセッションAPIがどこに位置づけられるかを示します。その後、小学校レベルの数学問題(GSM8K)を対象とした最小限のRLループを分解して、各オペレーションがどのように組み合わさるかを示し、最後にクックブックからの実世界のワークロードとして、競技プログラミングの問題に対する強化学習を取り上げます。ここでは、生成されたコードがユニットテストに対して実行され、その結果によって評価されます。
SageMaker AI トレーニングセッション API とは何ですか?
SageMaker AIのトレーニングセッションAPIを使うと、ローカルマシン上の自分のPythonプロセスから、稼働中のGPUバックエンドのセッションを駆動できます。操作を呼び出すと、SageMaker AIが、選択したベースモデルとLow-Rank Adaptation（LoRA）アダプターをロードしたマネージドアクセラレーター上でそれらを実行します。ジョブ全体を事前に記述する設定ファイルも、クラスター作成ステップもありません。各ステップで何が起こるかをコードから対話的に制御できるため、素早く反復やデバッグが行えます。
責任の分担は以下のようになっています。
| あなたが注目するのは | あなたが書く | SageMaker AI は処理します |
| あなたのデータ、RL環境、エージェント、アルゴリズム | undergone operations such as を呼び出す Python のトレーニングループ forward_backward(), optim_step()、および sample() GPUを必要とする手順について |
GPU割り当て、モデルおよびLoRAのロード、分散実行、障害復旧、チェックポイント保存、セキュリティ |
このAPIは、以下の3つのクライアントを中心に構成されています。 SageMaker Python SDK、順に作成されます。A ServiceClient プラットフォームレベルの設定(実行ロール、モデルパッケージグループ、Amazon Simple Storage Service (Amazon S3) の出力先、オプションで AWS Key Management Service (AWS KMS) キーと仮想プライベートクラウド (VPC) の設定。これらはセッションごとに上書きすることもできます)を保持し、セッションを開始します。The ServiceClient インスタンス化します TrainingClient (「trainingおよびcheckpointing操作など、 forward, forward_backward, optim_step、および save_weights_for_sampler)、これは続いて次をインスタンス化します: SamplingClient (sample および compute_logprobs、ロールアウトおよび評価中の生成とスコアリング用)。各API呼び出しはループ内の1つの操作に対応します。アクティブなセッションはOpenAI互換のchat completionsエンドポイントも提供するため、既存のOpenAIクライアントコードからサンプリングできます。
Figure 1: トレーニングセッション API における SageMaker Python SDK クライアントの階層構造
にとって 強化学習、ループには特徴的な形状があります。次の図は1回のトレーニングステップと、各部分がどこで実行されるかを示しています。緑色のステップはサービス管理のGPU上でAPIを呼び出し、紫色のステップであるループ制御と報酬ロジックはあなたのマシン上で実行されます。
図2: 1回のトレーニングステップ。サンプリング、勾配、チェックポイントの操作は管理対象GPU上で、ループ制御と報酬ロジックはあなたのマシン上で実行される
各ステップではプロンプトのバッチを読み込み、(1) GPU上の現在のモデルからサンプリングします。次にあなたのPythonコードがそれらの出力を(2) スコアリングします。検証器、報酬モデル呼び出し、テストハーネス——コードで表現できるものなら何でも構いません。(3) forward_backward および optim_step GPU上でスコアを勾配更新に変換し、(4) save_weights_and_get_sampling_client() 次のステップで更新後のモデルからサンプリングできるよう、サンプラーをリフレッシュします。
データ、ループ制御、報酬はユーザーが所有し、GPUに依存するサンプリング、勾配、チェックポイント処理はSageMaker AIが実行します。同じプリミティブで以下もカバーできます。 教師ありファインチューニング(SFT)およびその他のLLM学習技術. SFTの場合は、サンプリングと報酬のステップを省略し、 forward_backward と cross_entropy ラベル付き例に対する損失。
これは誰向けのもので、Amazon SageMaker AIの中でどのように位置づけられますか?
トレーニングセッションAPIは、カスタム損失関数、新しいRL報酬シグナル、マルチステージパイプラインなど、高速なイテレーションサイクルとアルゴリズムの自由度が必要で、かつGPUインフラを自分でプロビジョニングしたり運用したりしたくない場合に、最適な出発点です。
使用量は処理されたトークン単位で課金されるため、ループが消費した計算分のみが請求されます。報酬関数があなたのマシン上でバッチのスコアリングに10秒かかる場合、その実行中にGPU時間は課金されません。時間単位のGPU課金と比較して、これには2つの実用的な利点があります。第一に、アイドル状態のGPUコストが発生しないことです。報酬計算、データ読み込み、ロギング、デバッグは課金時間外で行われ、GPUに依存する操作を呼び出したときにのみ課金されます。第二に、コストが稼働時間やクラスターサイズではなく、データセットと生成量に応じて拡大するため、支出の予測が容易になり、実験の比較もしやすくなります。複雑な報酬モデル、マルチツールのエージェント型ロールアウト、人間参加型の評価など、ローカルでの重い計算を伴うワークロードでは、セッション全体でGPUを予約する場合と比べて節約額が大幅になる可能性があります。
トレーニングセッションAPIは、データセットだけでなく、トレーニングアルゴリズムによって仕事が定義されるチーム向けに構築されています。ループを自身で組み立てるため、同じプリミティブで、aを用いた教師ありファインチューニング（SFT）などに対応できます。 cross_entropy ラベル付きデータに対する損失、 強化学習 と共に ppo, importance_samplingまたは cispo (GRPO、PPO、およびその他のグループ相対的手法)、さらに カスタム目標選好最適化や蒸留など、同じ構成要素から作られる手法。以下に適合する:
- RL手法(GRPO、PPO、importance-samplingバリアント)、カスタムのadvantage estimator、または新しい目的関数を実装するポストトレーニングおよびアライメントの研究者。
- 設定駆動のレシピでは表現できない非標準の損失関数、データカリキュラム、あるいは多段階パイプラインを必要とする、教師ありファインチューニングや選好最適化を行うチーム向けです。
- カスタムロールアウトロジック、報酬設計（リワードシェーピング）、カリキュラムを必要とする、エージェント型およびマルチターンRLの開発者向けです。
- 設定駆動のファインチューニングを使いこなせるようになったが、GPUクラスタの運用負荷は引き受けたくない機械学習(ML)エンジニア。
- 低レベルのAPIに対して書かれたループに、操作ベースのプログラミングモデルが直接対応するため、既存の学習ループのコードを移植する研究者にとっても有用です。
forward_backward,optim_stepおよびsampleAPI。
サポートされているベースモデルと、それぞれが利用可能なAWSリージョンは、以下に記載されています。 トレーニングセッション API ドキュメント.
トレーニングセッションAPIがAmazon SageMaker AIの中でどう位置づけられるか
SageMaker AIのモデルカスタマイズオプションの中で、トレーニングセッションAPIは際立った位置にあります。フルマネージド型よりも、トレーニングアルゴリズムに対する大幅に高いコントロールを提供します。 サーバーレスモデルのカスタマイズ、そして Amazon SageMaker Training Jobs や Amazon SageMaker HyperPod.
| オプション | トレーニングの制御性 | インフラストラクチャ |
| SageMaker Serverless Model Customization – マネージドAPI | 低 – モデルを選択し、利用可能なハイパーパラメータを設定し、データを提供するだけ | サーバーレス、インスタンスの選択不要 |
| SageMaker Serverless Model Customization – SageMaker AI トレーニングセッションAPI | 高 – ループを自分で書き、プリミティブを呼び出す | サーバーレス、インスタンスの選択不要 |
| SageMaker Training Jobs | 完全 | マネージドの一時的ジョブで、インスタンスタイプを選択可能 |
| SageMaker HyperPod | 完全 | 高度な障害復旧機能を備えたマネージドの永続クラスター |
図3: トレーニングセッションAPIがSageMaker AIのモデルカスタマイズオプションの中でどこに位置するか
トレーニングループの構造: 数学問題におけるGRPO
このセクションでは、最小限の強化学習ループを分解して、各オペレーションがどのように組み合わさるかを示します。タスクは GSM8Kで、これは約8,500問の小学校レベルの数学文章題からなるベンチマークであり、各問題には次の例のような一つの数値解答が付いています:
Nataliaは4月に友人48人にクリップを売り、その後5月にはその半分の数のクリップを売りました。Nataliaが4月と5月に合わせて売ったクリップは何個ですか? → 72
OpenAIの GPT-OSS 20B 向けにGroup Relative Policy Optimization(GRPO)でLoRAアダプターを学習します。各質問に対してモデルは複数の応答を生成し、各応答は正しさでスコアリングされ、グループの平均を上回る応答が強化されます。完全に実行可能なサンプルは Amazon SageMaker examplesリポジトリ.
にあります。セッションを開く前に、ドキュメントの 前提条件 を完了してください。
セッションを開く
セッションを作成し、そこからトークナイザーと SamplingClient を取得します:
from sagemaker.train.training_session import ServiceClient, SamplingParams, AdamParams
svc = ServiceClient(role_arn=role, region=REGION)
training_client = svc.create_lora_training_client(
"openai/gpt-oss-20b", # a name from svc.list_base_models()
training_session_name=SESSION_NAME,
rank=32, alpha=64, # LoRA rank and scaling factor
model_package_group_arn=MODEL_PACKAGE_GROUP_ARN,
s3_output_path=OUTPUT_S3_PATH,
)
tokenizer = training_client.get_tokenizer()
sampling_client = training_client.create_sampling_client()報酬とアドバンテージを定義する
タスク固有のロジックは、あなたのマシン上で実行される2つの短い関数です:
def compute_reward(response_text, ground_truth):
predicted = extract_answer(response_text) # parse the number after "#### "
return 1.0 if predicted == ground_truth else 0.0
def compute_advantages(rewards):
rewards = np.array(rewards)
advantages = rewards - rewards.mean() # center within the group (the GRPO baseline)
return advantages / (rewards.std() + 1e-8) # normalizeループ
各ステップは図2に示す4つの段階に従います。 draw, build_promptと build_training_datum はサンプル内の短いヘルパーです:
for step in range(NUM_STEPS):
datums, pending = [], []
# (1) sample K responses per question in one request each (GPU); returns a future right away
for question, ground_truth in draw(PROMPTS_PER_STEP): # data loading is local
prompt = tokenizer.encode(build_prompt(question), add_special_tokens=False)
pending.append((prompt, ground_truth, sampling_client.sample(
prompt, num_samples=K,
sampling_params=SamplingParams(max_tokens=512, temperature=0.7, top_p=0.9))))
for prompt, ground_truth, future in pending:
result = future.result() # wait once the batch is in flight
# (2) score each response, then center the rewards within the group (local)
rewards = [compute_reward(tokenizer.decode(s.tokens), ground_truth) for s in result.sequences]
advantages = compute_advantages(rewards)
datums += [build_training_datum(prompt, s, a) for s, a in zip(result.sequences, advantages)]
# (3) one gradient step over the batch (GPU); the optimizer step is sequenced behind it
fwd_bwd_op = training_client.forward_backward(datums, loss_fn="ppo")
optim_op = training_client.optim_step(AdamParams(learning_rate=1e-5))
fwd_bwd_op.result(), optim_op.result()
# (4) refresh the sampler so the next step's rollouts come from the updated policy (GPU)
sampling_client = training_client.save_weights_and_get_sampling_client()
training_client.stop()このループを機能させている3つの詳細があります:
操作はfutureを返します。GPUに依存するすべての呼び出しは即座に APIFuture を返します。リクエストはキューに入り、バックグラウンドのサブミッターによって送信され、 .result() が値を待ちます。ループはどのリクエストも待つ前に、ステップのすべてのサンプリングリクエストをサブミットするため、バッチ全体が並行して生成を行います。学習操作は呼び出した順序で実行されるため、 optim_step は forward_backward が終了する前にサブミットできます。
データムは損失に必要なすべてを運びます。 build_training_datum はプロンプトと応答のトークン、 sample()からのトークンごとのlogprob、アドバンテージ、およびプロンプトを損失から除外する weights マスクをパックします。
チェックポイントはモデルパッケージです。すべての保存は、上記のループ内のステップごとのサンプラー更新を含め、モデルパッケージグループにバージョンを登録します。学習後、さらに2回の保存により、デプロイ用のマージ済みモデルと、重みとオプティマイザーを含む再開可能な状態が追加されます。
training_client.save_weights_for_sampler(merge_weights=True).result() # merged, deployable model
training_client.save_state().result() # resumable state (weights + optimizer)再開するには、最新のチェックポイントから新しいセッションを次のように開始します: create_training_client_from_state_with_optimizer.
実行中のセッションは、 stop() を呼び出すか期限が切れるまで、同時実行セッションのクォータを消費します。そのため、エラー経路を含めて終了時には必ず停止するか、クライアントを with ブロックで開いてください。マージされたチェックポイントはリアルタイムエンドポイントにデプロイできます。チェックポイントをLoRAアダプターとして提供する場合は、クックブックの デプロイ用ノートブック.
サンプルの実行では、40問に対する10回の学習ステップで、ホールドアウト精度が56パーセントから75パーセントに向上しました。
実世界のワークロード: 競技プログラミングに対するRL
GSM8Kでは仕組みが見えやすいままですが、本番のRLワークロードにはさらに多くの要素が加わります: より長い生成、マルチターンの対話、外部の採点プログラム、そしてGPUを稼働させ続けるのに十分なスループット。 SageMaker training sessions cookbook は、それらのエンジニアリングを、同じオペレーション上に構築された再利用可能な学習ループとしてパッケージ化しています。新しいRLタスクでは、環境を記述します: モデルが見るもの、その回答がどのように採点されるか、そして別のターンが与えられるかどうか。クックブックは、バッチ処理、チャットテンプレート、アドバンテージ推定、オーバーラップしたサンプリングと学習、検証、再開可能なチェックポイント、ステップごとのメトリクスを処理し、各実行の設定を記録して再現できるようにします。 code_rl レシピはこれを用いて、競技プログラミングの問題についてモデルを学習させます。
タスク。問題はDeepCoderデータセットから取得します（学習にはPrimeIntellect、TACO、LiveCodeBench、テストにはCodeforcesとLiveCodeBench）。モデルがPythonの解答を書き、その解答が問題のユニットテストに対して、お使いのマシン上の SandboxFusion コンテナで実行されます。報酬は、すべてのテストに合格すれば1、そうでなければ0で、応答にコードブロックが含まれない場合は0.1のペナルティが課されます。
マルチターンのツール使用。2ターン構成にすることで、モデルは check_solution ツールを呼び出し、テスト結果を読み、最終採点の前に回答を修正できます。ツールの呼び出しと結果はモデルのネイティブなチャット形式でレンダリングされるため、モデルは推論時に使用するのと同じトークン形式で学習されます。
このレシピが最小限のループに追加するもの:
- オーバーラップしたサンプリングと学習: ロールアウトワーカーは、トレーナーが1ステップを実行している間も生成を続けます。古いポリシーからサンプリングされた応答のグループは、最大kステップまでの古さであれば使用され、そうでなければ再キューイングされ、
cispo損失がポリシーの差を補正します。 - ストリーミングされる勾配計算:
forward_backwardは、ステップごとに1回ではなく、グループの完了に応じてサブバッチとして送信されます。 - ホールドアウト検証、チェックポイントからの再開、ステップごとのメトリクスは、オプションで次にも公開できます: SageMaker managed MLflow.
図4: Code RLのセットアップ。左側にローカルの学習ループとサンドボックス、右側にマネージドGPU
GPT-OSS 20B向けのLoRAアダプター（ランク32）を50ステップ学習しました。各ステップで128問、問題ごとに8つの応答を使用し、2ターン構成、staleness上限はk=3としました。学習報酬は0.46から0.60に上昇し、200問のホールドアウトテスト問題に対する精度は40パーセントから50.5パーセントに向上しました。
実行するには、サンドボックスを起動し、レシピを実行します:
bash tools/sandbox_up.sh
python -m training_sessions_cookbook.recipes.code_rl.train \
--config training_sessions_cookbook/recipes/code_rl/configs/coderl_oss20b_service.json \
--max_steps_off_policy 3 --stale_policy requeue --num_rollout_workers 80 --max_steps 50 \
--base_model_arn "$BASE_MODEL_ARN" \
--model_package_group_arn "$MODEL_PACKAGE_GROUP_ARN" \
--role_arn "$ROLE_ARN" \
--output_s3_path "$OUTPUT_S3_PATH"このクックブックには、同じ構造に関するレシピがさらに掲載されています。チャットデータセット(Tulu 3、UltraChat)を用いた教師ありファインチューニング、GSM8K でのフルスケールの強化学習(RL)、バイナリフィードバックからのKTO選好アライメントなどです。各レシピのREADMEには、コマンド、ハイパーパラメータ、期待されるメトリクスが記載されています。
まとめ
SageMaker AI トレーニングセッション API は、抽象化の境界をトレーニング操作のレベルまで引き下げます。自分のマシン上の Python でトレーニングループを記述すると、サービスがマネージド GPU 上でそれを実行し、モデルと LoRA のロード、分散実行、障害復旧、チェックポイント処理を担います。完全な GRPO ループは 4 つの操作に帰着します。 sample, forward_backward, optim_stepと save_weights_and_get_sampling_clientであり、報酬とアドバンテージのロジックは自分のコードに置かれます。Code RL レシピが示すように、同じ操作はサンドボックス化された採点を用いたマルチターンのツール使用 RL にも拡張できます。
カスタム報酬、エージェント的なロールアウト、ステップごとのデバッグ、新たに公開された手法は、すべて同じプログラミングモデルを使用します。アルゴリズムは自分のループに留まり、SageMaker AI が GPU を実行します。
始めるには:
- 次のレシピを選択してください: SageMaker training sessions cookbookのレシピ、例えば Code RL.
- Amazon SageMaker examples リポジトリの GSM8K example を実行し、自分の報酬関数に差し替えてください。
- 次のドキュメントをお読みください: training session APIs documentation 前提条件、クォータ、リージョン、サポートされるモデル、API リファレンスについてはこちらを参照してください。
謝辞
本記事への貢献とご指導をいただいた Arun Nagarajan、Xichao Wang、Shruti Sharma の各氏に感謝いたします。