新媒网 · 跨境数据社区

在数字经济浪潮中，我们跨境人始终走在探索前沿，尤其是AI大模型技术，正深刻改变着我们的业务模式和效率。然而，当我们面对日益庞大的AI模型时，一个核心挑战便浮现出来：如何高效地训练和优化这些“巨无霸”模型？这不仅是技术问题，更是决定我们能否抓住市场机遇的关键。

过去，我们很多朋友在进行强化学习训练时，习惯使用基于PyTorch DTensor（也就是FSDP2）的方案。这种方案确实为我们提供了一个方便的起点，它能很好地与HuggingFace生态系统融合，让初步尝试和扩展变得相对容易。无论是FSDP2、张量并行，还是序列并行、上下文并行，PyTorch的原生并行化能力都为我们提供了不少便利。然而，当模型规模膨胀到千亿级参数时，DTensor方案的局限性就显现出来了。巨大的激活内存消耗会带来显著的重新计算开销，导致训练步长（step time）变得难以忍受的慢，效率大打折扣。更重要的是，DTensor方案在NVIDIA CUDA核心优化和吞吐量提升方面，还缺少一些关键的“硬核”性能加持。

面对这些实际的痛点，一个更高效的解决方案显得尤为迫切。这正是NVIDIA Megatron-Core库诞生的使命所在。它专为解决超大规模语言模型的训练难题而设计，旨在提供更卓越的性能。

新媒网跨境获悉，NVIDIA NeMo-RL在最新的v0.3版本中，带来了对Megatron-Core后端支持的重大升级。这无疑为我们打开了一扇新的大门，让高效地对大型模型进行后训练成为可能。您可以在详细文档、示例脚本和配置文件中找到所需的全部信息，开始您的实战。

Megatron后端下的强化学习新篇章

Megatron-Core库的强大之处，在于它采用了针对GPU深度优化的技术和一系列高吞吐量性能增强方案。这让大规模语言模型的无缝训练成为现实。它独创的6D并行策略，能够精妙地优化通信和计算模式，并且支持各种复杂的模型架构。如今，NeMo-RL集成了对Megatron-Core的支持，这意味着开发者们在进行模型后训练时，也能充分利用这些顶级的优化技术。

或许有朋友会觉得，Megatron-Core提供了许多底层配置选项，对于初学者来说，这可能会让人感到有些“望而生畏”。但导师我想说，好的工具就是要让复杂变简单。NeMo-RL正是这样一款产品，它在幕后自动处理了大量的复杂调优工作，为用户呈现出一套更简洁、更直观的配置选项，让您能够更专注于模型本身。

迈出Megatron训练的第一步

启用基于Megatron的训练过程其实非常直接。您只需在您的YAML配置文件中，加入policy.megatron_cfg这一小节，就像是为您的模型插上了一对“性能之翼”。

policy:
 ...
 megatron_cfg:
  enabled: true
  activation_checkpointing: false
  tensor_model_parallel_size: 1
  pipeline_model_parallel_size: 1
 ...
 optimizer:
 ...
 scheduler:
 ...
 distributed_data_parallel_config:
  grad_reduce_in_fp32: false
  overlap_grad_reduce: true
  overlap_param_gather: true
  average_in_collective: true
  use_custom_fsdp: false
  data_parallel_sharding_strategy: "optim_grads_params"

配置文件中的所有参数都将在训练期间被转发给Megatron处理。当您将这一节加入配置并设置enabled=True后，您的模型便做好了训练的准备。启动训练的方式与使用DTensor时别无二致，您可以参考官方的README文档或我们关于DeepScaleR复现指南。

实战效果：数据不会说谎

基于Megatron的训练方案，同时支持密集型模型和混合专家（MoE）模型。下面的表格清晰展示了在使用GRPO（Group Relative Policy Optimization）对一些常用模型进行训练时的步长（step time）细分情况。表格中的时间数据，取自每个训练运行第22步至第29步的平均值，具有很强的参考价值。

表1：Megatron与PyTorch DTensor在不同训练配置下的模型性能对比

所有的运行都采用了相同的标准设置：最大序列长度4096，回滚批处理大小2048，全局批处理大小512，并启用了序列打包功能（关于序列打包的细节，我们稍后会提到）。针对Megatron-Core的运行，Llama 3.1 8B仅采用了数据并行。Llama 3.1 70B则采用了4路张量并行和4路流水线并行。Qwen3 32B采用了4路张量并行和2路流水线并行，而Qwen3 30B-A3B则使用了8路专家并行和2路张量并行。

值得注意的是，Llama 70B在DTensor模式下的结果，是使用动态批处理而非序列打包获得的，这是因为序列打包在该模式下存在已知的内存溢出问题。而Qwen3 32B和30B-A3B的DTensor运行均因断言错误而失败，这进一步凸显了Megatron-Core的稳定性优势。

通过充分利用Megatron-Core提供的性能优化，我们实现了远超DTensor的训练性能，同时，关键的模型收敛特性也得到了良好的保持。这意味着我们不仅跑得更快，而且模型的学习效果丝毫不受影响。

生成这些奖励曲线的命令如下，它们体现了在不同节点和并行设置下，如何配置运行：

8B模型 – 需单个节点：

## dtensor
uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B.yaml \
loss_fn.use_importance_sampling_correction=True

## megatron
uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B_megatron.yaml \
policy.sequence_packing.enabled=True loss_fn.use_importance_sampling_correction=True

70B模型 – 需8个节点：

## dtensor
uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_8B.yaml \
policy.model_name=meta-llama/Llama-3.1-70B policy.tokenizer.name=meta-llama/Llama-3.1-70B-Instruct \
policy.generation.vllm_cfg.tensor_parallel_size=4 policy.max_total_sequence_length=4096 \
cluster.num_nodes=8 policy.dtensor_cfg.enabled=True policy.dtensor_cfg.tensor_parallel_size=8 \
policy.dtensor_cfg.sequence_parallel=True policy.dtensor_cfg.activation_checkpointing=False \
loss_fn.use_importance_sampling_correction=True

## megatron
uv run ./examples/run_grpo_math.py --config examples/configs/grpo_math_70B_megatron.yaml \
policy.model_name=meta-llama/Llama-3.1-70B policy.tokenizer.name=meta-llama/Llama-3.3-70B-Instruct \
policy.sequence_packing.enabled=True loss_fn.use_importance_sampling_correction=True

这些运行方案都包含了性能和收敛性的增强配置，确保我们能在实现最优吞吐量的同时，也能保持模型的优异收敛表现。

1. 序列打包（Sequence packing）：这项技术能将多个序列“打包”到max_total_sequence_length中。它能有效减少填充令牌的数量，尤其当序列长度差异较大时，其效果尤为显著。以Llama 70B为例，启用序列打包能将总步长耗时缩短大约一倍，而且对模型收敛性没有任何负面影响。这项增强功能同时支持Megatron-Core和DTensor后端。想了解NeMo-RL中序列打包的更多细节，您可以查阅我们的文档。

2. 重要性采样（Importance sampling）：NeMo-RL为了追求最佳性能，在推理和训练中采用了不同的框架。然而，这可能会导致训练和推理的令牌概率之间存在细微差异。为了解决这个问题，一种有效的方法就是使用重要性采样。它会根据推理和训练概率的函数，为每个样本赋予一个权重。启用重要性采样可以减少不同运行之间的方差，让Megatron-Core和DTensor策略在收敛性上能更好地匹配。关于NeMo-RL中重要性采样的更多信息，同样可以在我们的文档中找到。

拓宽视野：长序列支持能力

借助Megatron-Core和DTensor，我们还能实现上下文并行，这对于长上下文训练至关重要。例如，下表展示了使用Megatron后端，在16k序列长度下，Llama 3.3 70B模型的当前性能表现。未来，更长的序列长度也将得到支持，并且长上下文训练的性能优化工作正在持续进行中。这无疑为我们处理更复杂、信息量更大的任务提供了强大的技术支撑。

表2：Llama 3.3-70B Instruct模型在16K长上下文窗口下，使用Megatron后端时的性能表现

其他值得关注的亮点功能

除了集成Megatron训练后端，NeMo-RL v0.3还引入了几项令人振奋的功能，这些功能让高效的后训练技术，能够惠及更广泛的模型，让更多开发者从中受益：

• 异步回滚（Async rollouts）：现在，用户可以通过设置policy.generation.async_engine=True来开启vLLM异步引擎，这将使多轮强化学习的训练速度提升2到3倍。这意味着您能够更快地迭代和优化模型，抢占市场先机。
• 非共置生成（Non-colocated generation） (DTensor后端)：用户现在可以选择将训练后端和生成后端放置在不同的GPU集合上。这项功能非常实用，特别是当训练和生成的并行模式/全局大小不兼容时，或者当在共置模式下，训练或生成的内存卸载后仍不足以降低内存占用时。这为我们提供了更大的灵活性，可以根据实际资源情况进行优化配置。更多细节，请查阅0.3.0版本的发布说明。

展望未来：即将上线的功能

未来总是充满无限可能。以下是即将上线的一些功能，值得我们持续关注和期待：

• 高效、更大规模的MoE模型支持：利用Megatron后端运行千亿级参数的模型，包括DeepSeek-V3和Qwen3-235B-A22B。这将把我们AI大模型的应用推向新的高度。
• 高度优化的修正（refit）：进一步提升模型修正的效率和性能。
• FP8生成支持：带来更快的推理速度和更低的内存消耗。
• Megatron和DTensor VLM支持：拓展对视觉语言模型（VLM）的兼容性。
• Megatron-Core后端下的非共置生成：为Megatron-Core用户提供更灵活的资源配置选项。

总结

在今天的分享中，我们深入探讨了NeMo-RL v0.3与Megatron-Core后端如何协同工作，显著提升了强化学习训练的吞吐量，尤其对于Llama 70B这类大型模型而言，效果尤为突出。凭借GPU优化的核心、6D并行策略，以及序列打包、重要性采样等实用功能，NeMo-RL确保了模型在不同规模下都能实现高效且稳定的收敛。我们还看到，它对长上下文训练提供了强大的支持，即便在16k的序列长度下，依然能保持优异的性能。

AI大模型时代的浪潮已至，效率是我们的生命线。希望今天的分享能为您的跨境事业注入新的活力。我鼓励大家积极探索NVIDIA NeMo RL的文档、示例配置和脚本，亲手实践，用Megatron-Core的优化能力，为您的AI大模型后训练提速，共同迎接更广阔的未来！

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

本文来源：新媒网 https://nmedialink.com/posts/17593.html