@PyTorch: 基于PyTorch、Ray、SGLang和NVIDIA Megatron-LM构建,Miles是RadixArk推出的一个用于大规模……的开源框架

X AI KOLs Timeline 工具

摘要

Miles是RadixArk推出的一个开源框架,用于大规模LLM强化学习后训练,集成了PyTorch、Ray、SGLang和NVIDIA Megatron-LM,支持MoE、低精度和容错。

基于PyTorch、Ray、SGLang和NVIDIA Megatron-LM构建,Miles是RadixArk推出的一个开源框架,用于大规模LLM强化学习后训练。 Miles使用PyTorch进行模型、数值计算、性能分析和扩展;使用Ray进行编排;使用SGLang进行rollout生成;使用Megatron-LM进行分布式训练。 该框架支持异步rollout和训练、NCCL/RDMA权重同步、MoE感知的rollout/训练对齐、低精度方案、LoRA、容错、可观测性,以及自定义算法和模型架构的扩展点。 阅读更多内容,请参阅Miles团队的最新博客:https://pytorch.org/blog/miles-a-pytorch-native-stack-for-large-scale-llm-rl-post-training/…
查看原文
查看缓存全文

缓存时间: 2026/06/30 21:47

基于PyTorch、Ray、SGLang和NVIDIA Megatron-LM构建,Miles是RadixArk推出的开源框架,用于大规模LLM强化学习后训练。

Miles使用PyTorch处理模型、数值计算、性能分析和可扩展性;Ray负责编排;SGLang用于生成推理结果;Megatron-LM则负责分布式训练。

该框架支持异步推理与训练、NCCL/RDMA权重同步、MoE感知的推理/训练对齐、低精度方案、LoRA、容错机制、可观测性,并为自定义算法和模型架构提供扩展点。

了解更多信息,请阅读Miles团队的最新博客:https://pytorch.org/blog/miles-a-pytorch-native-stack-for-large-scale-llm-rl-post-training/…


面向大规模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/)

摘要

Miles 是 RadixArk 推出的开源框架,用于大规模 LLM RL 后训练。它整合了 SGLang(用于推理生成)、NVIDIA Megatron-LM(用于训练)、Ray 编排以及 PyTorch 原生的可扩展性,并内置了统一的低精度方案、MoE 感知的推理/训练对齐、快速的 NVIDIA NCCL/RDMA 权重同步、可观测性和容错机制——使前沿规模的 LLM 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:训练器 ranks、SGLang 推理服务器、路由代理以及异步推理工作节点都位于 Ray 的 Actor 模型内。

这为 Miles 提供了集群级 RL 工作负载的自然基础。

将工作节点放置在 GPU 上

Miles 使用 Ray 的 GPU 感知调度器和放置组进行 Actor 放置,通过启动时的 Ray 放置规格支持分离式(推理和训练在不同节点上)和同地部署(推理和训练在同一节点上)布局。进程放置必须考虑机架感知,以促进精细的同地部署、预留备用节点,这对于错误隔离至关重要,因为隔离机架内的问题(例如,区分坏 GPU 与整个机架问题)并不总是简单明了。

在 RL 流水线中移动数据

提示、样本和更新后的权重在推理 Actor 和训练器 ranks 之间持续循环,Miles 使用 Ray Actor 和任务来协调这一流程。对于批量权重传输,Ray 处理控制路径,而张量字节则通过专用的 NCCL/RDMA 通道传输,这使得 Miles 兼具 Ray 级别的可编程性和大数据量的快速路径。

监控长时间运行的作业

由于 Miles 运行是一个端到端的 Ray 作业,它继承了 Ray 的操作界面——作业提交、工作节点监控、日志聚合和仪表盘可见性——无需额外的基础设施。启用容错机制后,Miles 可以恢复失败的 ranks,并在相同的 Ray 基础上保持持续数周的工作负载运行。

支持完全异步的 RL

因为 Ray Actor 是持久的、拥有自己的状态,并且独立调度,Miles 可以运行完全异步的模式,其中推理和训练不再相互阻塞——推理 Actor 持续将样本流式传输到队列中,训练器按自己的节奏从队列中取出数据。

Megatron-LM:扩展训练后端

Miles 将 Megatron-LM 作为其生产训练后端,直接接入 Megatron 的参数解析器、模型构建流水线、训练循环、并行原语和分布式检查点格式,而不是将其作为黑盒库进行封装。这使得 Miles 拥有前沿规模密集和 MoE 训练所需的基础设施,同时保持用户友好的工作流程。

统一的参数接口

Megatron-LM 已经暴露了大型分布式训练的配置接口——序列长度、旋转位置编码、分组 GEMM、各种并行方式、优化器设置、激活检查点等——Miles 直接复用这些接口,而不是封装或重新声明。用户通过一个启动脚本配置 Miles 运行,该脚本结合了 Miles 特定选项和标准 Megatron 选项,避免了重复的配置层,并使训练设置接近上游 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 后训练更易于复现、扩展和运维。

相似文章