`torchforge`破局！大模型RL训练简化，驾驭700亿参数。

当下，强化学习（RL）已成为前沿人工智能发展的核心要素，其应用范围涵盖指令遵循、推理以及复杂的科研任务。通过反馈机制提升模型性能，已超越传统静态数据集学习模式，转变为模型通过尝试行动并接收奖励来优化行为。然而，实践中，强化学习的实际研究工作往往被其基础设施的复杂性所困扰。

经典的强化学习训练流程通常包括：智能体（Agent）在环境中采取行动，获得奖励，并根据奖励更新其行为策略（Policy），即决定行动的策略。在大语言模型（LLM）场景中，策略本身即是模型，更新策略意味着调整模型参数以提升响应质量。尽管理论上看似简单，实际操作中却面临分布式基础设施的诸多挑战：

规模问题：许多现代模型规模庞大，无法单机容纳（有时甚至无法单主机容纳），因此需要进行权重分片。不同工作负载（如训练和推理）常需采用不同的分片策略。
性能瓶颈：在面向大语言模型的强化学习中，奖励通常来源于策略模型的输出。这要求进行自回归生成，成本较高。若训练过程完全依赖生成环节而阻塞，将严重影响吞吐量，从而拖慢训练速度。
工具与环境：在可验证奖励的强化学习（RLVR，奖励来源于代码执行或数学检查等客观验证）中，通常需要对策略模型生成的结果进行正确性验证（例如，检查数学计算是否正确，或代码能否通过单元测试）。此外，模型可能需要在训练过程中学习使用各种工具。这两种情况下，不同工具的延迟可能显著影响系统的端到端性能。
权重同步：在分布式部署中，当训练生成新的策略权重时，这些权重需要同步传输至所有推理副本。对于一个拥有700亿参数的模型，若分布在16个副本上，这意味着需要传输数百GB的模型权重。传统的网络传输方式在此情境下速度过慢，一次更新可能耗时数分钟，完全成为训练迭代速度的瓶颈。

为解决上述性能瓶颈，通常需要允许一定程度的“离线策略”（Off-policyness），即允许轨迹生成循环基于略微过时的权重进行。这就催生了权重存储抽象的必要性，它与回放缓冲区（Replay Buffer）共同作为训练工作器和策略工作器之间的中介。仅仅异步化不足以解决问题。不同工作器可能拥有不同的性能特征，因此能够根据需要灵活扩展或缩减不同组件以消除瓶颈，对于实现峰值性能至关重要。现有框架常要求显式的移动、重分片和复杂的控制流来管理大规模异步强化学习循环。

新媒网跨境了解到，Meta近期发布了torchforge。torchforge的推出旨在解决这些问题，其目标是抽象化基础设施的复杂性，使研究人员能够专注于算法本身。torchforge的设计理念主要体现在以下几个方面：

以伪代码方式表达强化学习逻辑，同时能无缝扩展至GPU集群。
支持任意程度的异步性，从完全同步的PPO到完全异步的离线策略训练。
将基础设施与算法分离，让研究人员专注于核心算法开发。
自然组合，能够混合搭配组件以适应不同的强化学习方法。
基于成熟基础：构建于PyTorch原生的分布式框架Monarch（用于协调和容错）、Meta生产级大语言模型训练平台torchtitan以及高吞吐量、内存高效的推理引擎vLLM之上。

本篇文章将深入探讨torchforge的核心功能，包括其可组合的API，Monarch如何抽象化基础设施复杂性并实现可扩展性，以及vLLM和torchtitan等生产级组件如何从实际部署中提供性能和验证。

需要说明的是，Monarch和torchforge目前仍处于实验阶段，并持续积极开发中。其API可能会随着早期用户的反馈和方法的完善而发生变化。

化强化学习为伪代码

在torchforge中，生成单个训练情节（episode）的逻辑示例展现了其简洁性：

async def # 从数据集中采样一个提示 # 使用策略模型（vLLM推理）生成响应 # 计算参考模型对KL惩罚的对数概率 await  # 评估响应质量 await # 存储情节用于训练 await

这部分代码直接反映了强化学习的核心逻辑——采样提示、生成响应、计算奖励、为KL散度惩罚计算参考对数概率。这正是算法设计时在伪代码中可能出现的表述。值得注意的是，这些代码中并未包含：

重试逻辑或故障处理。
资源分配或版本管理。
同步逻辑。

而当我们将这段逻辑组合成连续的异步训练流程时，其简洁性依然保持：

async def int """异步强化学习：连续生成轨迹，连续训练。"""  # 启动多个并发轨迹生成器  for _ in range # 启动连续训练  await async def  """使用最新策略持续生成轨迹。""" while await async def   """持续训练可用的经验数据。""" 0 while  # 从回放缓冲区中采样当前训练版本的经验数据 await if is None 0 1 experience else await 1   # 推送更新后的权重并广播到策略副本 await await

多个轨迹生成器可以并行运行（通过num_rollout_loops参数调整）。一旦经验数据可用，训练便会持续进行。这是一种完全异步的离策略强化学习模式，能够实现最大吞吐量。

异步训练有助于提升吞吐量，但如果需要实现PPO或其他要求严格数据收集保证的在策略（on-policy）算法，情况则有所不同。generate_episode()函数仍然适用，只需通过不同的组合方式即可：

async def synchronous_rl(batch_size: int ):  """同步在策略强化学习：收集批次数据，然后训练。"""  version = 0   while True:  # 使用当前策略版本收集完整的批次数据   for _ in range (batch_size):    await generate_episode(     dataloader, policy, reference_model,     reward, replay_buffer    )   # 采样刚刚收集的批次数据   batch = await replay_buffer.sample.route(    curr_policy_version=version,    batch_size=batch_size   )   # 对完整的批次数据进行训练   inputs, targets = batch   loss = await trainer.train_step.route(inputs, targets)  # 推送更新后的权重并广播到策略   await trainer.push_weights.route(version + 1 )   await policy.update_weights.fanout(version + 1 )   version += 1

在此模式下，协调模式发生了变化。系统会在训练前收集一个完整的批次数据，且权重更新与数据收集同步进行。但轨迹生成逻辑和强化学习代码保持不变。这正是torchforge的优势所在：研究人员可以编写一次轨迹生成逻辑，并将其组合到任何范式中，无论是 ऑन-policy、off-policy 还是介于两者之间。

这种简洁性并非没有代价。在这些API背后，是一个复杂的、处理容错、权重同步、分布式协调等的基础设施层。这些复杂性并未渗入到用户的代码中，但在底层默默工作，使得这种体验成为可能。那么，它是如何实现的呢？

基础支撑：Monarch

torchforge在很大程度上依赖于Monarch，这是一个基于可扩展Actor消息传递机制的PyTorch原生分布式编程框架。Monarch提供了一个统一的控制器编程模型，使得分布式强化学习工作负载易于管理。

SPMD问题

在分布式机器学习领域，通常将工作负载表达为SPMD（Single Program, Multiple Data）作业。这种编程范式下，控制逻辑是在副本层面定义的。在一个典型的面向大语言模型的强化学习训练设置中，包含生成工作器（负责生成文本）、训练工作器（更新权重）、奖励评估器（评估输出）以及协调这些组件的回放缓冲区。各个组件间的连接代表数据移动、协调或同步，这些都需要进行管理。

在实际部署中，模型常因规模过大无法完全置于单个GPU上，这意味着生成器和训练器需要跨多个设备进行分片。由于这些工作负载通常以SPMD作业形式存在，即控制逻辑在副本级别定义，因此对这些组件间的数据流进行建模变得复杂。在SPMD编程模型中，数据和控制流难以理解。研究人员必须同时考虑多个进程，协调集合操作以实现快速数据移动，并管理不同进程间通信的复杂性。

Monarch的方法

得益于Monarch以Mesh为中心、基于Actor的消息传递模型，分布式工作负载能够被优雅地建模。

Monarch并未要求研究人员关注单个进程和集合操作，而是允许在逻辑组件层面进行编程。生成器、训练器和奖励模型都被视为ActorMesh（Actor集合）。组件间的通信可以自然地表达，例如“生成器，对这个提示进行推理”或“训练器，这里有一个批次数据用于训练”。Monarch的单一控制器负责协调所有操作。用户只需编写一个Python程序来协调逻辑，调用Actor的方法并传递数据。控制器会处理将这些逻辑操作映射到分布式执行的复杂性，包括哪些进程进行通信、如何协调集合操作以及如何进行RDMA传输。

重要的是，底层的组件本身通常是SPMD作业——vLLM用于推理时采用张量并行，torchtitan用于训练时采用FSDP。Monarch并非取代这些组件，而是对其进行编排。这使得集成现有已在大规模场景下表现良好的生态系统组件变得容易。用户既能享受到成熟SPMD实现带来的性能和可扩展性优势，又能获得更简洁的编程模型（单一控制器协调）。

关于Monarch如何实现大规模扩展并同时提供容错性和快速数据移动，其中涉及大量的分布式系统设计决策。对技术细节感兴趣的读者可以查阅Monarch相关博客文章。对于torchforge而言，其主要基于Monarch的轻量级服务抽象（构建于分布式Actor之上）以及TorchStore（一个高效的、分布式的、基于RDMA数据传输的内存键值存储）进行构建。

服务（Services）

torchforge引入了“服务”这一更高级别的抽象，它构建在Monarch的Actor之上，负责管理分布式Actor的所有操作复杂性：包括在节点间生成副本、容错、负载均衡以及智能路由。在torchforge中，组件可以作为服务创建，并附带简单的资源规格：

# 创建一个策略服务，包含16个副本，每个副本使用8个GPU
policy = PolicyActor.options(
  hosts= 1 ,
  procs= 8 ,
  with_gpus=True,
  num_replicas= 16 ).as_service()
# 创建一个轻量级编码环境服务
# 用于在编码任务的强化学习中，安全地执行生成的代码
coder = SandboxedCoder.options(
 procs= 1 ,
 with_gpus=False,
  num_replicas= 16 ).as_service()
# 创建一个奖励模型服务
reward = RewardActor.options(
  hosts= 1 ,
  procs= 4 ,
 with_gpus=True,
  num_replicas= 4 ).as_service()

.options() API允许用户指定每个服务的形态和规模。例如，上述策略服务创建了16个副本，每个副本横跨8个GPU，总共为推理提供了128个GPU，并实现了副本间的自动负载均衡。编码环境则是一种不同类型的服务——轻量级、仅使用CPU，并随作业临时创建。无需单独的Kubernetes部署。用户可以启动16个副本以在轨迹生成期间并行执行代码，并在作业完成后自动关闭。

服务修饰词（Service Adverbs）

服务提供了在副本级别操作的基础修饰词：

# route() - 将请求负载均衡到单个副本
response = await policy.generate.route(prompt)
result = await coding_env.execute.route(code_string)
# fanout() - 广播到所有副本
await policy.update_weights.fanout(version)

这些修饰词在强化学习代码中随处可见。当调用policy.generate.route(prompt)时，torchforge会执行以下操作：

选择一个可用的策略副本（例如通过轮询方式）。
将请求路由到该副本的Actor。
返回响应。
自动避开故障副本。

有状态操作

当需要保持一致性时，可使用粘性会话（sticky sessions）：

async with policy.session():
 # 在此上下文中，所有调用都命中同一个副本
  # 对于在不同回合中保持KV缓存状态至关重要
  response1 = await policy.generate.route(prompt1)
  response2 = await policy.generate.route(prompt2)

容错能力

如果副本在轨迹生成过程中发生故障，服务会检测到并将其标记为不健康，然后将后续请求路由到健康的副本。故障副本会自动重启。强化学习代码不会感知到这些故障。

对强化学习的重要性

这些服务抽象解决了强化学习基础设施中的关键挑战，这些挑战通常迫使开发人员编写分布式系统代码而非强化学习算法：

轨迹生成间的负载均衡：在异步强化学习中，多个continuous_rollouts()任务并发运行，每个任务都在生成情节。服务会自动将这些轨迹分布到可用的副本上。无需手动负载均衡逻辑，无需管理工作器池——只需调用.route()，torchforge便会处理副本选择。
异构扩展：不同组件需要不同资源。策略服务可能需要16个副本x8个GPU以实现高吞吐量vLLM推理。奖励模型可能需要4个副本x4个GPU。编码环境可能需要16个轻量级、仅CPU的副本。服务允许每个组件根据其瓶颈独立扩展。
临时基础设施：服务随作业创建，并在完成后销毁。无论是尝试新的奖励模型，还是在两次运行之间添加新的环境，都像修改Python代码一样简单。无需维护常驻部署，也无需提前预置额外基础设施。
容错确保训练持续运行：强化学习训练可能持续数小时或数天。硬件故障在所难免。如果没有服务级别的容错机制，单个GPU故障意味着整个作业需要重启并丢失进度。有了服务，故障副本会自动重启，健康副本则持续处理轨迹。训练得以不间断进行。
强化学习代码中无基础设施：回顾generate_episode()和训练循环的代码，其中不包含工作器管理、重试逻辑、故障处理或负载均衡代码。服务在基础设施层处理所有这些问题，从而使强化学习算法保持简洁。

服务解决了控制平面问题——请求路由、副本管理、故障处理。但分布式强化学习还有另一个关键挑战：数据平面。每个训练步骤都会产生新的检查点——可能包含数百亿参数，需要从训练器传输到所有策略副本。对于运行vLLM推理的16个策略副本，需要将一个700亿参数的模型广播16次。如果传输缓慢，权重同步将成为瓶颈。如果操作不当，可能导致检查点损坏或内存不足。这就是TorchStore发挥作用的地方。

TorchStore：数据平面

面对权重同步的挑战，开发者通常有两种合理的选择：

设计一个通过互连的P2P操作网络解决重分片问题的系统。这种方法需要管理复杂的分布式集合网络，不仅繁琐，而且难以通用化。此外，用户还必须在离策略训练情况下同时管理内存中额外的检查点副本，同时考虑具有有限DRAM/VRAM的不同分布式拓扑结构。
结合分布式检查点（DCP）和网络文件系统（NFS）。这种用户体验是最好的情况——用户只需调用“保存”和“加载”，无需关心传输细节或不同分布式拓扑下的重分片，因为这些大多交由NFS处理。但实际中，网络文件系统通常并非设计用作传输缓冲区，并且这种方法存在巨大的基础设施成本和性能瓶颈。

新媒网跨境了解到，Torchstore的诞生源于一个简单的观察：我们应该能够结合使用中央存储（如NFS）的用户体验和内存中P2P操作的性能。鉴于Monarch的出现，我们现在也拥有了设计此类系统的合适原语。具体来说，TorchStore是一个分布式的内存键值存储，专为PyTorch张量优化，并构建于Monarch原语之上。

Torchstore致力于为用户提供以下特性：

最佳用户体验：利用简单的DTensor API，在该特定领域提供最佳的用户体验。
最佳性能：尽可能接近硬件成本的极限性能。
完全灵活的存储安排：用户可以根据需要存储张量——与训练器/生成器协同放置、拥有独立的存储层、分片/不分片——并且只需少量代码修改即可更改存储方式。

以下是其在实践中的应用：

import torchstore as ts
from torch.distributed._tensor import distribute_tensor, Replicate, Shard
from torch.distributed.device_mesh import init_device_mesh
async def place_dtensor_in_store():
  device_mesh = init_device_mesh( "cpu" , ( 4 ,), ...)
  tensor = torch.arange( 4 )
  dtensor = distribute_tensor( tensor, device_mesh, placements=[Shard( 1 )]
 )
  # 存储该进程的dtensor分片
 await ts.put( "my_tensor" , dtensor)
async def fetch_dtensor_from_store()
  device_mesh = init_device_mesh( "cpu" , ( 2 , 2 ), ...)
  tensor = torch.rand( 4 )
  dtensor = distribute_tensor(
   tensor,
   device_mesh,
   placements=[Replicate(), Shard( 0 )]
  )
  # 使用本地dtensor信息获取正确位置的数据
 await ts.get( "my_tensor" , dtensor)
if __name__ == "__main__" :
  ts.initialize()
  run_in_parallel(place_dtensor_in_store, world_size= 4 )
  run_in_parallel(fetch_dtensor_from_store, world_size= 4 )
  ts.shutdown()

实际上，torchforge通过集成TorchStore立即获得了两大优势。原生的DTensor支持意味着无论训练器端的权重如何分片，在生成过程中都可以请求任意张量切片，这封装了大部分重分片逻辑，避免了在torchforge中进行定制化集成。此外，该集成利用大量可用的CPU内存来存储权重副本，使GPU能够不间断运行。这使得训练和生成完全解耦，消除了训练循环中昂贵的训练器/生成器同步需求，这是实现完全异步工作负载的重要一步。TorchStore对分布式存储的通用支持也使其成为强化学习数据平面中其他多种用例的首选。例如，未来的优化可以利用Torchstore原语来实现生成器输出的不同回放策略，或作为更大规模容错解决方案的一部分。

强化学习栈：成熟组件与简洁协调

torchforge在设计时有意识地避免了重复造轮子。它没有构建自己的推理引擎或训练框架，因为这些领域已经存在许多出色的、经过实战检验的解决方案。可以将torchforge理解为高效组合成熟组件的框架。vLLM负责高吞吐量推理，torchtitan管理可扩展训练，而定制环境则执行代码或与工具交互。这些组件带来了效率、可扩展性和经过验证的性能。torchforge的作用是协调：使这些组件无缝协作，让研究人员能够自然地表达其强化学习算法。组件承担繁重的工作，torchforge负责编排，从而保持强化学习代码的简洁。

集成vLLM与TorchTitan

torchforge当前集成了vLLM用于推理，以及TorchTitan用于训练。这些集成正在不断演进——团队正致力于更清晰的边界和上游贡献，以使Monarch在这些框架中获得原生支持。vLLM提供PagedAttention、连续批处理以及经过大规模验证的吞吐量。torchforge直接与vLLM的引擎集成，让用户能够根据研究需求定制生成策略、内存管理或推理逻辑。TorchTitan作为生产级的训练基础设施，提供FSDP、流水线并行、张量并行以及经过大规模验证的优化。其集成允许直接访问训练步骤逻辑和分片策略，在不受框架限制的情况下进行实验。

目前，这些集成允许直接访问组件代码。这既是其特色（可深度定制），也是一个正在进行中的工作（团队正在完善模式）。随着这些集成的稳定，团队将分享构建自定义组件的最佳实践。核心思想依然是——torchforge在逻辑层面协调组件。无论是像vLLM和TorchTitan这样成熟的框架，还是用户自行构建的定制实现，torchforge都负责编排，让用户专注于强化学习算法。

可扩展环境

强化学习通常需要与文本生成之外的环境进行交互，例如执行代码、使用工具、运行模拟。torchforge通过相同的服务抽象将这些环境视为一等公民。其首个环境集成是沙盒代码执行。对于编码任务的强化学习，用户需要安全地执行生成的代码并评估结果：

# 用于并行执行的轻量级纯CPU服务
coder = SandboxedPythonCoder.options(
  procs= 1 , with_gpus=False, num_replicas= 16 ).as_service()

随后在强化学习代码中：

async def generate_episode():
 prompt = await dataloader.sample.route()
 code = await policy.generate.route(prompt)
 stdout, stderr = await coder.execute.route(code)
 reward = stderr == "" # 例如，如果代码成功运行则为1
  await replay_buffer.add.route(Episode(...))

服务使得环境具有临时性——随作业启动，独立扩展，并在完成后拆卸。与策略和奖励模型一样，服务的相同协调原语也适用于环境。

迈向智能体工作流

这些可扩展环境是迈向智能体（Agentic）强化学习工作流的基础性步骤。智能体需要与工具交互、执行代码、查询API并导航复杂的环境，同时从结果中学习。编码环境展示了其模式：将Python组件封装在Actor中，将其转换为服务，然后它便能自然地与强化学习栈组合。同样的方法也适用于其他环境类型——包括如HuggingFace的Open Environments等新兴框架，这些框架旨在标准化工具和环境接口。

这仅仅是一个开始。随着强化学习向更具智能体能力的方向发展，团队期待与社区合作，构建下一代环境集成。

合作伙伴与验证案例

新媒网跨境获悉，torchforge项目获得了斯坦福大学Scaling Intelligence Lab和Coreweave的宝贵支持。斯坦福大学团队将其弱验证器项目Weaver集成到torchforge中，使得模型能够针对MATH和GPQA等具有挑战性的推理基准进行训练。CoreWeave提供了强大的512块H100 GPU集群，支持了torchforge的大规模训练运行，为实验提供了流畅高效的体验。

后续规划

torchforge正处于发展初期。团队已经构建了基础——可组合的原语、简洁的API和成熟组件的集成——但仍有许多工作有待完成。

磨平棱角：早期使用者可能会遇到一些问题，例如文档不完善、错误信息不明确、API不一致等。团队承诺根据实际使用反馈进行快速迭代。
扩展环境支持：编码沙盒是torchforge的首个环境集成。团队正与HuggingFace和Prime Intellect合作开展Open Environments倡议，以标准化工具和环境接口，从而更方便地将多样化环境接入强化学习训练工作流。
扩展至专家混合模型（MoE）：目前的工作侧重于密集模型，但torchforge的设计考虑了支持大规模MoE模型。团队将继续努力，提升torchforge的规模和性能。

如果研究人员在构建强化学习系统时，基础设施的复杂性阻碍了研究进展，可以尝试torchforge。

新媒网（公号: 新媒网跨境发布），是一个专业的跨境电商、游戏、支付、贸易和广告社区平台，为百万跨境人传递最新的海外淘金精准资讯情报。

本文来源：新媒网 https://nmedialink.com/posts/torchforge-simplifies-rl-training-70b-params.html