
SageMaker训练会话让开发者用Python控制循环,由服务管理GPU执行
据AWS介绍,Amazon SageMaker AI训练会话API允许开发者用Python编写自己的训练循环,同时由服务承担GPU执行工作。开发者可以调用forward_backward、optim_step和sample等操作,逐步组织模型更新与生成过程。这些调用作用于持续存活、保留状态的会话,模型在不同步骤之间保持加载在GPU上。
服务负责的工作包括运行环境设置、分布式执行、梯度同步、故障恢复和租户隔离。开发者则通过自己的Python代码安排训练步骤。这里的控制方式以会话及其操作为核心:训练循环可以包含生成、评分、更新和再次生成,而基础运行工作交由服务处理。
从本地Python控制训练
用户可以从本地Python进程驱动会话,所选基础模型及LoRA适配器加载在托管加速器上。按照AWS描述,这种编程模型无需先创建集群,也无需预先编写一份描述整个任务的配置文件。循环中的具体调用由Python程序组织,模型则持续驻留在托管计算资源中。
SDK按顺序使用三个客户端。ServiceClient保存平台设置并启动会话;TrainingClient提供训练操作和检查点操作;SamplingClient提供sample及compute_logprobs。活跃会话还提供兼容OpenAI的聊天补全端点。三个客户端分别对应平台与会话启动、训练与保存,以及生成与概率计算。
文章展示的强化学习循环先让托管GPU生成回答,再由本地Python代码计算评分。随后,程序调用forward_backward和optim_step更新模型,并通过save_weights_and_get_sampling_client刷新采样客户端。由此,生成、局部评分和参数更新构成了可在Python中编排的循环。
监督微调使用相同的基础操作,但省去采样和奖励计算,直接对带标签样本使用cross_entropy损失。AWS还列出ppo、importance_sampling和cispo等强化学习损失,以及偏好优化、蒸馏等自定义目标。不同目标的区别体现在训练循环和损失的组织方式中。
异步调用、顺序与计费
与GPU执行有关的调用会立即返回APIFuture对象,请求在后台排队并提交。程序调用result()时等待对应操作完成。采样请求可以一起提交,以并发生成回答;训练操作则按照提交顺序执行。因此,文章明确区分了生成请求的并发方式与训练更新的执行顺序。
示例中的训练数据不仅包含提示和回答的token,还包括采样时记录的逐token对数概率、优势值,以及排除提示token损失的掩码。这些字段把生成阶段的信息、回答的相对评价和参与损失计算的位置放在同一条训练记录中。提示内容虽然存在于记录里,但通过掩码不计入该示例的损失。
计费方面,AWS表示,费用依据GPU相关操作处理的token数量计算。本地奖励计算、数据加载、日志记录或调试期间,不收取GPU时间费用。材料没有给出每token的具体价格,也没有提供可用于量化成本节省的数据,因此这里能够确认的是所描述的计费边界。
AWS将训练会话定位在几种服务模式之间:相较完全托管的Serverless Model Customization,它提供更多算法控制;相较SageMaker Training Jobs或HyperPod,用户承担的运维工作更少。这个定位与文章的编程方式相呼应,即用户组织训练算法,服务负责运行和分布式执行。
检查点、恢复与配方
每次保存检查点都会注册一个模型包版本。示例分别保存用于部署的合并模型,以及包含权重和优化器状态、用于继续训练的状态。恢复训练使用create_training_client_from_state_with_optimizer。这两类保存结果分别服务于部署和恢复训练,示例对它们作了区分。
合并检查点可以部署到实时端点。会话在停止或到期之前,都会占用并发会话配额;保存检查点与结束会话是材料中分别描述的事项。给定资料没有列出并发配额的具体数值,因此无法从示例推导可同时运行多少个会话。
配套训练会话配方集提供可复用的监督微调和强化学习流程,并处理批次组织、聊天模板、优势估计、采样与训练重叠、验证及检查点恢复。它还记录每一步的指标和运行配置。对于阅读示例的开发者,这些配方展示了基础API如何被组织为更完整的训练过程。
列出的配方包括使用Tulu 3和UltraChat的聊天监督微调、完整规模的GSM8K强化学习,以及根据二元反馈进行KTO偏好对齐。材料同时提供数学任务和编程任务的实例,展示生成结果如何进入评分,再由评分参与模型更新。
数学示例的结果由AWS报告
GSM8K示例使用GRPO训练GPT-OSS 20B的LoRA适配器。其奖励逻辑关注组内相对表现:正确性得分超过同组平均水平的回答会得到强化。AWS报告,在40道题上执行10个训练步骤后,留出测试数据上的准确率由56%升至75%。
这些数字描述的是文章给出的特定样例运行,关联着指定模型、数据及训练步骤。它们是AWS报告的结果,没有提供独立验证,也不能据此把相同提升幅度推广到其他模型、数据集或训练任务。
编程配方支持测试与修改
code_rl配方使用DeepCoder中的竞赛编程题,将生成的Python解答放入本地SandboxFusion容器,以单元测试评分。全部测试通过时奖励为1,否则为0;回答缺少代码块时,还会受到0.1的惩罚。评分规则因而同时包含测试结果和代码块要求。
该配方采用两轮设置,允许模型调用check_solution、查看测试结果,再修改自己的答案。测试反馈由此进入同一次两轮交互:模型可以在得到检查结果后修订解答,而后续训练使用这一流程生成的回答。
在执行安排上,编程配方让回答生成与训练重叠进行,同时限制回答的陈旧程度。它使用cispo校正策略差异,并以子批次提交梯度计算。生成与更新重叠、限制滞后以及校正策略差异,是材料共同描述的训练安排。
AWS报告的编程运行使用秩为32的GPT-OSS 20B LoRA适配器,训练50步,涉及128道题,每题生成8个回答,采用两轮交互,并将陈旧程度限制在3个训练步骤以内。在该运行中,训练奖励从0.46升至0.60;在200道留出题目上的准确率从40%升至50.5%。指标也可选择发布到SageMaker托管MLflow。
上述编程指标同样只对应AWS列出的具体运行。现有材料没有枚举支持的基础模型、可用AWS区域、数值配额或每token价格。因而,这份介绍能够说明会话编程方式、训练配方和报告中的实验结果,但不能补出这些未给出的可用性与价格细节。