NVIDIA NeMo-RL集成Megatron-Core，赋能百亿级模型高效后训练

NVIDIA NeMo-RL v0.3引入了对Megatron-Core库的支持，旨在解决PyTorch DTensor（FSDP2）在训练百亿参数以上的大型模型时遇到的性能瓶颈。

Megatron-Core通过GPU优化内核、6D并行策略以及高效的通信和计算模式，显著提升了训练吞吐量。与DTensor相比，Megatron-Core在处理大规模模型时，能够更有效地管理激活内存，减少重计算开销，从而大幅缩短训练步长时间。例如，在Llama 3.1 70B模型的训练中，Megatron-Core后端将总步时间从230秒缩短至147秒，提升了约56%。

NeMo-RL通过自动化复杂配置，简化了Megatron-Core的集成过程，使得开发者能够更便捷地利用其性能优势进行模型后训练。该优化支持密集模型和混合专家（MoE）模型，并能有效处理长序列训练（如16K上下文长度），同时保持与DTensor相当的收敛性。此外，新版本还增加了异步Rollouts和非共置生成等功能，进一步提升了训练效率和灵活性。

Reinforcement Learning with NVIDIA NeMo-RL: Megatron-Core Support for Optimized Training Throughput | NVIDIA Technical Blog

The initial release of NVIDIA NeMo-RL included training support through PyTorch DTensor (otherwise known as FSDP2). This backend enables native integration with the HuggingFace ecosystem…

NVIDIA Technical BlogAnna Shors