跳到正文
AWS Machine Learning Blog· Nicolas Jourdan·· 4 小時前精選AI 評分66

Amazon SageMaker AI training session APIs 如何讓用戶自訂訓練迴圈而毋須管理 GPU

Write your own training loop — without managing GPUs — with Amazon SageMaker training session APIs

AI 導讀

AWS 為 Amazon SageMaker AI 加入 training session APIs，讓用戶以本機 Python 自訂訓練迴圈，並由服務管理 GPU 運算。

推薦理由

文章以 GSM8K 數學題及競賽編程 RL 範例，拆解本機訓練迴圈如何組合抽樣、獎勵計算與梯度更新，並由 SageMaker 管理 GPU 執行及檢查點。

正文 · 繁體中文

現時微調大型語言模型（LLM），通常要在訓練迴圈控制權與營運簡便性之間作取捨。使用全代管服務時，你只需選擇基礎模型、提供資料和超參數，其餘一切均由服務負責；當標準的預先封裝方案適合你的問題時，這種方式十分有效。管理程度較低的服務則給你更多控制權，但你也要管理更多底層基礎設施和訓練框架。有時標準方案並不足夠：你可能需要為監督式微調（SFT）自訂損失函數或資料課程、自訂偏好最佳化方法、強化學習（RL）獎勵函數，或上星期才發表的對齊演算法。一旦需要自行設計訓練程序，就得承擔相關的基礎設施工作，以換取演算法上的自由。

今日，Amazon SageMaker AI 的無伺服器模型自訂服務新增了訓練工作階段 API，將抽象層級下移一層，消除了這種取捨。它們不再把演算法抽象化，而是把運算資源抽象化。forward_backward()、optim_step() 和 sample() 等操作會轉化為 SageMaker API 呼叫，你可以按所用技術的需要，組合成任何程序：SFT、偏好最佳化、RL，或全新的方法。這些呼叫會指向一個正在運作且具狀態的工作階段，因此模型會在各個步驟之間繼續載入於 GPU 上，而你可以按需要向它發出任意數量的呼叫。服務會處理執行環境設定、分散式執行、梯度同步、故障復原、租戶隔離，以及大規模執行 GPU 的運作。你編寫訓練迴圈，SageMaker AI 則負責執行基礎設施。

我們亦同步推出 SageMaker 訓練工作階段實作指南，內含一系列可直接執行的 SFT 和 RL 方案，你可以此作為藍本，編寫自己的訓練迴圈。

本文會介紹這種程式設計模型和三種 SDK 用戶端，並説明訓練工作階段 API 在 SageMaker AI 的模型訓練及自訂選項中所處的位置。接著，我們會拆解一個以小學數學題（GSM8K）為例的簡易 RL 迴圈，説明各項操作如何配合，最後介紹實作指南中的實際工作負載：針對競賽編程題進行強化學習，並透過執行單元測試為生成的程式碼評分。

甚麼是 SageMaker AI 訓練工作階段 API？

透過 SageMaker AI 訓練工作階段 API，你可以在本機電腦的 Python 程序中，操控一個即時運作、由 GPU 支援的工作階段。你呼叫相關操作，SageMaker AI 便會在受管理的加速器上執行，並載入你選定的基礎模型和低秩適配（LoRA）適配器。你無需預先使用設定檔描述整項工作，也無需建立叢集。你可以透過程式碼逐步控制執行內容，並即時互動，因此能快速反覆測試和偵錯。

雙方的職責分工如下：

由你負責 你編寫 SageMaker AI 負責
你的資料、RL 環境、Agent 和演算法 Python 訓練迴圈；在需要使用 GPU 的步驟中呼叫 forward_backward()、optim_step() 和 sample() 等操作 GPU 分配、模型和 LoRA 載入、分散式執行、故障復原、檢查點儲存、安全性

此 API 以 SageMaker Python SDK 中的三個用戶端為核心，並依序建立。ServiceClient 保存平台層級的設定（執行角色、模型套件羣組、Amazon Simple Storage Service (Amazon S3) 輸出位置，以及可選的 AWS Key Management Service (AWS KMS) 金鑰和虛擬私有雲 (VPC) 設定；你也可以按工作階段覆寫這些設定），並啟動工作階段。ServiceClient 會實例化 TrainingClient（訓練和檢查點操作，例如 forward、forward_backward、optim_step 和 save_weights_for_sampler），而後者再實例化 SamplingClient（sample 和 compute_logprobs，用於 rollout 和評估期間的生成及評分）。每次 API 呼叫都對應迴圈中的一項操作。啟用中的工作階段亦會提供相容於 OpenAI 的聊天完成端點，因此你可以使用現有的 OpenAI 用戶端程式碼，從中抽樣。

SageMaker Python SDK client hierarchy: ServiceClient creates a TrainingClient, which in turn creates a SamplingClient

圖 1：SageMaker Python SDK 訓練工作階段 API 的用戶端層級結構

對於強化學習，迴圈有其典型流程。下圖展示一個訓練步驟，以及各部分的執行位置：綠色步驟會呼叫服務管理 GPU 上的 API；紫色步驟則是你的迴圈控制和獎勵邏輯，會在你的機器上執行。

One reinforcement learning training step, showing which operations run on managed GPUs and which run on the local client

圖 2：一個訓練步驟；抽樣、梯度和檢查點操作在受管理的 GPU 上執行，而迴圈控制和獎勵邏輯則在你的機器上執行

每個步驟都會載入一批提示詞，並 (1) 在 GPU 上從目前的模型抽樣。接着，你的 Python 程式碼會 (2) 為這些輸出評分：可以是驗證器、獎勵模型呼叫、測試框架，或任何你能用程式碼表達的方式。(3) forward_backward 和 optim_step 會將評分轉換為 GPU 上的梯度更新，而 (4) save_weights_and_get_sampling_client() 會更新抽樣器，讓下一個步驟從更新後的模型抽樣。

你負責資料、迴圈控制和獎勵；SageMaker AI 則執行受 GPU 限制的抽樣、梯度和檢查點操作。同一套基本操作亦適用於監督式微調 (SFT) 及其他 LLM 訓練技術。進行 SFT 時，移除抽樣和獎勵步驟，並在標註範例上使用 cross_entropy 損失呼叫 forward_backward。

這適合哪些人？在 Amazon SageMaker AI 中又有甚麼定位？

當你需要快速迭代並保有演算法上的自由度，例如使用自訂損失、新的 RL 獎勵訊號或多階段流程，但又不想自行配置或操作 GPU 基礎設施時，訓練工作階段 API 是理想的起點。

費用按處理的 token 數量計算，因此你只需支付迴圈實際使用的運算資源。如果你的獎勵函數在自己的機器上需要 10 秒為一批資料評分，在這段時間內不會收取 GPU 費用。與按小時收取 GPU 費用相比，這有兩項實際優勢。第一，沒有閒置 GPU 成本：獎勵計算、資料載入、記錄和偵錯都在 GPU 計費時間以外進行，只有在呼叫受 GPU 限制的操作時才會收費。第二，成本會隨數據集和生成量增加，而非取決於實際經過時間或叢集大小，因此更容易預測支出，也更容易比較實驗。對於本機計算量大的工作負載，例如複雜的獎勵模型、多工具代理式 rollout，或人在迴圈中的評估，與整個工作階段都預留 GPU 相比，節省的費用可能相當可觀。

訓練工作階段 API 專為工作內容取決於訓練演算法，而不只是數據集的團隊而設。由於你可以自行編寫迴圈，相同的基本操作可用於監督式微調（SFT），以 cross_entropy 損失函數處理已標註數據；用 ppo、importance_sampling 或 cispo（GRPO、PPO 及其他羣組相對方法）進行強化學習；以及使用相同元件建立自訂目標，例如偏好最佳化或蒸餾。適用對象包括：

  • 實作 RL 方法（GRPO、PPO、重要性抽樣變體）、自訂優勢估算器或新目標的後訓練及對齊研究人員。
  • 進行監督式微調或偏好最佳化，並需要非標準損失函數、數據課程安排，或設定檔驅動配方無法表達的多階段管道的團隊。
  • 需要自訂 rollout 邏輯、獎勵塑形及課程安排的 Agent 型及多輪 RL 系統建構團隊。
  • 已超出設定檔驅動微調方式所能應付範圍，但不想承擔 GPU 叢集維運負擔的機器學習（ML）工程師。
  • 移植現有訓練迴圈程式碼的研究人員；以操作為基礎的程式設計模型，可直接對應以低階 forward_backward、optim_step 及 sample API 編寫的迴圈。

支援的基礎模型，以及各模型可用的 AWS 區域，列於訓練工作階段 API 説明文件。

訓練工作階段 API 在 Amazon SageMaker AI 中的定位

在 SageMaker AI 的模型自訂選項中，訓練工作階段 API 處於獨特的位置：與完全受管理的Serverless Model Customization相比，它讓你大幅增加對訓練演算法的控制；與在Amazon SageMaker Training Jobs或Amazon SageMaker HyperPod上訓練相比，則大幅減少你需要自行負責的維運工作。

選項 對訓練的控制 基礎架構
SageMaker Serverless Model Customization – 受管理 API 低 – 選擇模型、設定可用的超參數並提供數據 無伺服器，無需選擇執行個體
SageMaker Serverless Model Customization – SageMaker AI 訓練工作階段 API 高 – 自行編寫迴圈並呼叫基本操作 無伺服器，無需選擇執行個體
SageMaker Training Jobs 完整 受管理的暫時性任務，可自行選擇執行個體類型
SageMaker HyperPod 完整 受管理的持續運作叢集，具備進階故障復原功能

Chart plotting control over the training algorithm against infrastructure you manage, showing training session APIs at high control and low infrastructure relative to other SageMaker AI options

圖 3：訓練工作階段 API 在 SageMaker AI 模型自訂選項中的定位

訓練迴圈剖析：數學問題上的 GRPO

本節會拆解一個精簡的強化學習迴圈，展示各項操作如何配合。任務採用GSM8K，這個基準測試包含約 8,500 道小學數學應用題，每題只有一個數字答案，例如：

Natalia 在 4 月把夾子賣給 48 位朋友，到了 5 月，她賣出的夾子數量只有 4 月的一半。Natalia 在 4 月和 5 月總共賣出多少個夾子？ → 72

我們使用羣組相對策略最佳化（GRPO），為 OpenAI 的 GPT-OSS 20B 訓練 LoRA 適配器：模型會為每條問題生成多個回應，並按正確性為每個回應評分；表現高於羣組平均值的回應會獲得強化。完整且可執行的範例載於Amazon SageMaker 範例儲存庫。

開啟工作階段前，請先完成説明文件中的先決條件。

開啟工作階段

建立工作階段，然後從中取得分詞器和一個 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()

定義獎勵和優勢

任務專屬邏輯由兩個在本機執行的簡短函數組成：

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 所示的四個階段進行。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()

有三個細節令這個循環得以運作：

操作會傳回 future 物件。每個需要 GPU 的呼叫都會立即傳回 APIFuture：要求會排入佇列，由背景提交程序送出，而 .result() 會等待取得值。循環會先提交某個步驟的所有採樣要求，再等待任何一項結果，因此整個批次可以並行生成。訓練操作會依照呼叫順序執行，因此可在 forward_backward 尚未完成前提交 optim_step。

資料項目包含計算損失所需的一切。build_training_datum 會打包提示詞和回應 token、來自 sample() 的逐 token 對數機率、優勢值，以及 weights 遮罩，確保提示詞不會計入損失。

檢查點是模型套件。每次儲存都會在模型套件羣組中註冊一個版本，包括上述循環中每個步驟的採樣器更新。訓練完成後，再儲存兩次，加入供部署使用的合併模型，以及包含權重和最佳化器、可供恢復訓練的狀態。

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 訓練工作階段教學指南將這些工程實作整合成可重用的訓練循環，並以相同的操作為基礎。對於新的 RL 任務，你只需編寫一個環境，定義模型會看到甚麼、答案如何評分，以及模型能否再進行一輪。教學指南會處理批次處理、聊天範本、優勢估算、重疊採樣與訓練、驗證、可恢復的檢查點，以及每步指標，並記錄每次執行的設定，以便重複執行。code_rl 訓練方案會利用這套工具，訓練模型解答競技程式設計題目。

任務。題目來自 DeepCoder 數據集（訓練使用 PrimeIntellect、TACO 和 LiveCodeBench；測試使用 Codeforces 和 LiveCodeBench）。模型會編寫 Python 解答，並在本機的 SandboxFusion 容器中執行，對照題目的單元測試。所有測試通過時，獎勵為 1，否則為 0；如果回應不含程式碼區塊，則扣 0.1。

多輪工具使用。設定兩輪互動後，模型可以呼叫 check_solution 工具、讀取測試結果，並在最終評分前修改答案。工具呼叫及結果會以模型原生的聊天格式呈現，因此模型接受訓練時使用的 token 格式，與推理時使用的格式相同。

這個訓練方案在最精簡的循環上增加了以下功能：

  • 重疊進行採樣與訓練：訓練器執行訓練步驟時，生成工作程序會持續生成回應。若一組回應來自較舊的策略，且其舊度不超過 k 個步驟，就會使用該組回應；否則會重新排入佇列，而 cispo 損失會校正策略差異。
  • 串流式梯度計算：每當一組完成，forward_backward便會以子批次提交，而不是每個步驟只提交一次。
  • 留出集驗證、檢查點續接及逐步指標，並可選擇發佈至SageMaker 託管式 MLflow。

Code RL setup: a local training loop and SandboxFusion container on the left, and managed GPUs running GPT-OSS 20B with a LoRA adapter on the right

圖 4：Code RL 配置，左側為本機訓練迴圈和沙盒，右側為受管理的 GPU

我們使用兩輪互動及 k=3 的陳舊度上限，為 GPT-OSS 20B 訓練了一個 LoRA 配接器（秩為 32），共訓練 50 步，每步包含 128 個問題，每個問題有 8 個回應。訓練獎勵由 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 迴圈可簡化為四項操作：sample、forward_backward、optim_step及save_weights_and_get_sampling_client，獎勵和優勢計算邏輯則由你在自己的程式碼中處理。正如 Code RL 配方所示，相同的操作亦可擴展至使用沙盒評分的多輪、工具使用型 RL。

自訂獎勵、Agent 式 rollout、逐步偵錯及新近發表的技術，都採用相同的程式設計模型：演算法留在你的迴圈中，而 SageMaker AI 負責執行 GPU 工作。

開始使用：

致謝

感謝 Arun Nagarajan、Xichao Wang 和 Shruti Sharma 為本文作出的貢獻和提供的指導。

關於作者

Nicolas Jourdan

Nicolas Jourdan

Nicolas 是 AWS 資深 AI/ML 專業解決方案架構師，常駐德國法蘭克福。他協助客戶在 AWS 上建構 AI 和機器學習工作負載，涵蓋早期設計決策以至投入生產。他的專長包括模型自訂、分散式訓練及推理。加入 AWS 前，他曾任研究科學家及工程經理，從事自動駕駛、無人機和製造業相關應用，並持有物件偵測和異常偵測方面的專利。

Bruno Pistone

Bruno Pistone

Bruno 是 AWS 全球生成式 AI/ML 資深專業解決方案架構師，常駐意大利米蘭。他與 AWS 產品團隊及大型客戶合作，協助他們充分了解技術需求，並設計能充分發揮 AWS 雲端及 Amazon ML 技術堆疊優勢的 AI 和機器學習解決方案。他的專長包括分散式訓練及推理工作負載、模型自訂、生成式 AI，以及端對端機器學習。他喜歡與朋友共度時光、探索新地方，以及到訪新的目的地。

Giuseppe Angelo Porcelli

Giuseppe Angelo Porcelli

Giuseppe 是 Amazon Web Services 的首席機器學習專家解決方案架構師。憑藉多年軟件工程經驗及機器學習背景，他與不同規模的客戶合作，瞭解他們的業務及技術需求，並設計能充分運用 AWS 雲端及 Amazon Machine Learning 技術堆疊的 AI 和機器學習解決方案。他曾參與多個不同領域的項目，包括 MLOps、電腦視覺及 NLP，並運用多種 AWS 服務。閒暇時，Giuseppe 喜歡踢足球。

Pooja Karadgi

Pooja Karadgi

Pooja 是 AWS 的高級產品經理，負責領導 Amazon SageMaker AI 的模型自訂工作。她致力讓基礎模型更易取用，並切合生產環境的實際需要，協助客戶透過強化學習等技術，按其特定領域調整模型。她亦曾參與 SageMaker Studio 的開發者及管理員使用體驗工作，涵蓋從實驗到生產的整個機器學習生命週期。

來源：AWS Machine Learning Blog · aws.amazon.com

#產品更新#數據/訓練#部署/工程#Agent