Miles:用于大规模LLM强化学习后训练的PyTorch原生栈(14分钟阅读)

TLDR AI 工具

摘要

Miles是RadixArk推出的一个开源的PyTorch原生框架,用于大规模LLM强化学习后训练,集成了SGLang、Megatron-LM和Ray,以实现高吞吐量的推演和分布式训练。

Miles是一个用于大规模LLM强化学习后训练的框架。它让前沿规模的LLM强化学习更易于构建、复现和操作。随着模型变得越来越大,并在更加分布式的专用硬件上运行,强化学习后训练已经成为一个分布式系统问题。Miles使大规模LLM强化学习训练更具组合性、可复现性,且更易于扩展,同时保持核心训练器足够精简,以便研究人员和基础设施团队进行定制。
查看原文
查看缓存全文

缓存时间: 2026/07/01 17:19

# 一个用于大规模 LLM RL 后训练的 PyTorch 原生栈 – PyTorch 来源: https://pytorch.org/blog/miles-a-pytorch-native-stack-for-large-scale-llm-rl-post-training/ ### 特色项目 - PyTorch 标识 (https://pytorch.org/projects/pytorch/) - Ray 标识 (https://pytorch.org/projects/ray/) ## **TL;DR** Miles 是 RadixArk 开源的面向大规模 LLM RL 后训练的框架。它整合了用于推理生成的 SGLang、用于训练的 NVIDIA Megatron-LM、Ray 编排以及 PyTorch 原生可扩展性,并封装在一个小巧、可插拔的训练器中,内置统一的低精度方案、MoE 感知的推理/训练对齐、快速的 NVIDIA NCCL/RDMA 权重同步、可观测性和容错能力——使得构建、复现和运行前沿规模的大模型 RL 变得更加容易。 ## **为什么是 Miles?** 强化学习已成为大型语言模型后训练的核心部分。但随着模型变得更大,从密集型转向混合专家模型(MoE),并在更多分布式的专用硬件(如 NVIDIA Blackwell 和 Hopper 系列)上运行,RL 后训练不再仅仅是一个训练循环。它成了一个分布式系统问题。 一个现代的 LLM RL 框架需要协调多个移动部件: - 推理工作节点必须高吞吐量地生成样本。 - 训练器必须高效地消费这些样本并计算稳定的策略更新。 - 推理策略和训练策略必须保持同步。 - 大型 MoE 模型引入的路由行为必须在推理和训练之间保持一致。 - 低精度方案需要在全流程中一致工作。 - 长时间运行的任务从一开始就需要可观测性、检查点和容错能力。 Miles 正是为此场景而构建的。 Miles 是 RadixArk 开源的面向 LLM 后训练的强化学习框架。它原生构建于 SGLang 之上以实现高吞吐量推理,并与 Megatron-LM 深度集成以实现可扩展训练,使用 Ray 编排分布式系统,并在整个技术栈中保持 PyTorch 作为通用的编程和数值层。 目标很简单:让大规模 LLM RL 训练更具可组合性、可复现性,且更易于扩展,同时保持核心训练器足够小巧,以便研究者和基础设施团队进行定制。 ## **Miles 架构** Miles 遵循小而核心、多扩展点的理念。 核心训练循环故意保持紧凑。用户最常想改变的部件——推理逻辑、奖励计算、损失函数、样本过滤、指标和训练循环钩子——在启动时通过用户提供的 Python 模块附加。这使得团队能够在不分支整个框架的情况下,使系统适应新算法和生产约束。 在这个小核心之下,Miles 整合了四个主要系统: - **SGLang** 用于高吞吐量推理生成。 - **Megatron-LM** 用于可扩展的分布式训练。 - **Ray** 用于集群编排、actor 生命周期、调度和监控。 - **PyTorch** 用于模型、自动梯度、分布式原语、数据类型支持、扩展性和性能分析。 这种组合至关重要。RL 后训练要求生成和训练协同工作,但这两个阶段具有截然不同的性能特征:推理受内存带宽限制(解码过程中 KV 缓存和参数读取占主导),而训练受计算限制且通信密集。权重同步、样本传输、检查点转换、路由一致性和低精度行为都需要在边界上谨慎处理。 本文的其余部分将介绍 Miles 如何处理该边界的每个部分——使用 Ray 进行编排,使用 Megatron-LM 进行扩展,使用 PyTorch 进行扩展,以及开箱即用的功能。 ## **Ray:编排长时间运行的 RL 任务** Miles 直接构建在 Ray 分布式运行时之上。在 Miles 运行中,每个长期存在的进程都表示为 Ray actor:训练器 rank、SGLang 推理服务器、路由代理和异步推理工作节点都位于 Ray 的 actor 模型中。 这为 Miles 提供了集群规模 RL 工作负载的自然基础。 ### **将工作节点放置到 GPU 上** Miles 使用 Ray 的 GPU 感知调度器和放置组来实现 actor 放置,通过启动时的 Ray 放置规约支持**分离式**(推理和训练在不同节点上)和**同地部署**(推理和训练在相同节点上)两种布局。进程放置必须考虑机架感知,以便于仔细的同地部署、预留备用节点,并且对于错误隔离至关重要,因为隔离机架内的问题(例如,区分故障 GPU 和整个机架问题)并非总是直截了当。 ### **在 RL 流水线中移动数据** 提示、样本和更新后的权重在推理 actor 和训练器 rank 之间持续循环流动,Miles 使用 Ray actor 和任务来协调该流程。对于批量权重传输,Ray 处理控制路径,而张量字节通过专用的 NCCL/RDMA 通道移动,从而为 Miles 同时提供 Ray 级别的可编程性和大数据传输的快速路径。 ### **监控长时间运行的任务** 由于 Miles 运行是一个端到端的 Ray 任务,它继承了 Ray 的操作接口——任务提交、工作节点监控、日志聚合和仪表盘可见性——无需额外的基础设施。在启用容错功能的情况下,Miles 可以恢复故障 rank,并在相同的 Ray 基础上保持长达一周的工作负载持续运行。 ### **支持完全异步的 RL** 因为 Ray actor 是持久的,拥有自己的状态并独立调度,Miles 可以运行完全异步模式,其中推理和训练不再相互阻塞——推理 actor 持续将样本流式传输到队列中,训练器按自己的节奏从队列中消耗。 ## **Megatron-LM:扩展训练后端** Miles 使用 Megatron-LM 作为其生产训练后端,直接接入 Megatron 的参数解析器、模型构建流水线、训练循环、并行原语和分布式检查点格式,而非将其包装为黑盒库。这为 Miles 提供了前沿规模密集型和 MoE 训练所需的基础设施,同时保持用户友好的工作流程。 ### **统一的参数界面** Megatron-LM 已经暴露了大量分布式训练配置参数——序列长度、旋转位置嵌入、分组 GEMM、各种并行策略、优化器设置、激活检查点等——Miles 直接复用这些参数,而不是包装或重新声明。用户通过一个结合了 Miles 特有选项和标准 Megatron 选项的启动脚本来配置 Miles 运行,避免了重复的配置层,并使训练设置接近上游 Megatron 的行为。 ### **模型规范而非长期分支** 前沿架构变化迅速,新的注意力模块、路由机制和专家布局不断出现在各模型家族中,因此 Miles 通过可插拔的**模型规范**来处理它们——这些小型规范文件将自定义的 PyTorch 组件(例如,门控注意力输出模块、Gated-Delta-Net 块或特定模型的 MoE 路由器)直接插入 Megatron 的模型流水线。这使得 Miles 能够支持新架构——例如 DeepSeek-V3/V4、GLM-4.7 和 Qwen3 MoE 变体——而无需维护一个长期分支、不断与上游分歧的 Megatron 分支。 ### **并行感知的检查点** Miles 使用 Megatron 的并行感知分布式检查点格式,因此模型可以从 Hugging Face 转换一次,然后根据不同的张量/流水线/上下文/专家并行配置加载,无需从头重新转换权重。对于运营大型训练任务的团队来说,这意味着每当模型或集群形状发生变化时,检查点转换和并行性更改不会成为一个单独的工程项目。 ### **无需修补后端即可扩展训练** Miles 在训练循环中明确定义的点处暴露钩子——模型初始化后、对数概率计算前、以及每个训练步骤前——因此用户可以添加辅助损失、自定义指标、样本级诊断、裁剪规则或算法特定行为,而无需编辑 Megatron 内部实现。设计目标很简单:保持后端强大,但将用户定制放在后端之外。 ## **PyTorch:模型、数值和可扩展性的公共层** PyTorch 是 Miles 内部的通用编程模型:模型组件是常规的 `torch.nn.Modules`,损失是标准的自动梯度图,混合精度、梯度检查点、分布式原语和性能分析都保持在熟悉的 PyTorch 工作流中。这很重要,因为 LLM RL 后训练变化很快——团队需要添加新的奖励、损失、路由器、模型模块和调试工具,而无需每次都学习新的抽象。 ### **PyTorch 原生模型可扩展性** Miles 的可插拔模型规范机制围绕 `torch.nn.Modules` 构建,因此支持新架构意味着将新组件编写为普通的 PyTorch 代码,并将其连接到 Megatron 的模型流水线中——自动梯度、混合精度、梯度检查点和模块生命周期都按照 PyTorch 用户期望的方式工作。团队无需将模型转换为单独的中间抽象即可在 Miles 上运行。 ### **PyTorch 原生 RL 定制** 同样的原则适用于 RL 算法:推理函数、奖励、损失函数、样本过滤器、指标和训练循环钩子均通过启动时提供的 Python 模块进行定制,这些模块使用标准的 PyTorch 操作,并与训练图的其他部分组合。团队可以从现有配方开始,替换奖励、添加辅助损失、更改样本过滤或引入新的诊断,而无需重写训练器。 ### **跨流水线的低精度方案** Miles 在其低精度流水线中构建于 PyTorch 的数据类型系统之上,提供 BF16、FP8、MXFP8 和 INT4-QAT 方案,这些方案涵盖训练和推理,而不是作为孤立的仅后端功能。这种一致性对于 RL 很重要,因为用于生成样本的策略和用于计算训练对数概率的策略必须保持对齐,而 Miles 旨在使这些数值选择明确且可重复。 ### **使用熟悉工具进行性能分析和调试** 大规模 RL 性能问题可能出现在任何地方——推理延迟、训练计算、集体通信、数据移动、权重同步、样本过滤或调度——因此 Miles 集成了 PyTorch 性能分析器,以捕获训练阶段的 Chrome 跟踪,供标准工具检查。结合 Megatron 基于 PyTorch 的后端以及受支持情况下的图编译路径,这使调试和性能优化保持在熟悉的 PyTorch 生态系统内。 ## **Miles 开箱即用的功能** Miles 旨在提供大规模 LLM RL 后训练所需的核心系统功能: - **推理与训练集成**——连接 SGLang 推理与 Megatron-LM 训练,支持分离式和同地部署执行,以适应不同的 GPU 预算和利用率目标。 - **异步执行**——完全异步模式将推理与训练解耦:推理 actor 持续将样本流式传输到队列中,训练器按自己的节奏从队列中消耗,消除了两阶段之间每次迭代的阻塞。 - **快速权重同步**——每次训练更新后,新权重通过专用 NCCL/RDMA 通道流向推理工作节点,Ray 仅处理控制路径,因此批量张量字节不经过 Python 数据路径。 - **MoE 感知的推理/训练对齐**——推理路由回放保留跨推理/训练边界的路由决策,减少了否则会破坏 MoE RL 的训练器与推理路由不匹配问题。 - **低精度支持**——统一的 BF16 / FP8 / MXFP8 / INT4-QAT 流水线,作为端到端 RL 栈的一部分设计,而非孤立的仅训练方案。 - **跨推理和训练的 LoRA**——在推理和训练路径中都支持 LoRA,实现参数高效的后训练,降低大型基模型的成本并加速迭代。 - **容错性和可观测性**——Ray 的任务和 actor 模型提供监控、日志聚合和仪表盘可见性,而 rank 级容错保证长达一周的训练运行持续进行;PyTorch 性能分析器集成覆盖训练级视图。 - **广泛的模型和硬件支持**——Miles 为前沿和开源模型提供即用配方,包括 DeepSeek-V4、Kimi K2.5 / K2.6、GLM-5 / 5.1 和 Qwen3.5 / 3.6,并支持 NVIDIA 旗舰 Hopper / Blackwell GPU。 ## **小而核心,众多扩展点** Miles 最重要的设计选择之一是核心训练器保持小巧。 Miles 不是强迫用户为每种新算法或模型家族分支框架,而是暴露明确的扩展点: - **推理函数**用于自定义生成行为。 - **奖励函数**用于任务特定监督。 - **损失函数**用于新的 RL 目标。 - **样本过滤器**用于数据选择和拒绝。 - **训练钩子**用于指标、诊断、辅助损失和自定义更新逻辑。 - **模型规范**用于架构特定模块。 这些扩展点使 Miles 适用于广泛的后训练工作流:经典的 RLHF 式训练、基于规则的奖励训练、代码和智能体任务、MoE 后训练、低精度实验以及需要自定义可观测性或安全检查的生产流水线。 简而言之,Miles 处理系统级决策——放置、权重同步、容错、低精度方案——以便用户代码可以专注于算法和产品逻辑。 ## **展望未来** LLM 后训练发展迅速——更大的模型、更长的上下文、更多的 MoE,以及更异步、更智能体化、系统密集型的 RL 流水线——而 Miles 正是为这一轨迹而构建的:通过在一个小巧的可插拔训练器背后整合 SGLang、Ray、Megatron-LM 和 PyTorch,它为研究者和基础设施团队提供了一条从算法实验到大规模 RL 运行的 PyTorch 原生路径,这也是我们开源 Miles 以让前沿规模 LLM RL 后训练更易于复现、扩展和运营的原因。

相似文章