NVIDIA通过NeMo-RL的Megatron-Core提升训练吞吐量
realtime news Aug 20, 2025 14:18
NVIDIA在NeMo-RL v0.3中引入Megatron-Core支持，通过GPU优化技术和增强的并行性优化大模型的训练吞吐量。
NVIDIA推出了其NeMo-RL框架的最新版本0.3，其中包含对Megatron-Core的支持。据NVIDIA官方博客介绍，这一增强旨在通过利用GPU优化技术和先进的并行策略来优化大语言模型的训练吞吐量。
以往后端的挑战
NVIDIA NeMo-RL的初始版本使用了PyTorch DTensor (FSDP2)，提供了与HuggingFace生态系统的本地集成，并通过PyTorch的本地并行性支持快速实验。然而，随着模型规模增加到数百亿个参数，DTensor路径由于显著的重计算开销和缺乏优化的NVIDIA CUDA内核而显得不够，导致步骤时间低效。
引入Megatron-Core
Megatron-Core库通过提供更高效的大规模模型训练解决方案解决了这些限制。它采用6D并行策略来增强通信和计算模式，支持各种模型架构。这个后端实现了大规模语言模型的无缝训练，显著提升了吞吐量和性能。
开始使用Megatron-Core
实现基于Megatron的训练需要向YAML配置添加特定配置。NeMo-RL简化了这一过程，自动处理复杂的调优，为用户提供简便的配置选项。这使得开发者更容易采用Megatron-Core，让他们专注于优化模型训练过程。
性能提升
基于Megatron的训练支持密集模型和专家混合（MoE）模型。性能测试表明，与PyTorch DTensor相比，Megatron-Core在各种模型配置中（如Llama 3.1-8B和70B）表现出优越的训练性能，增强体现在更快的步骤时间和改进的收敛特性。
附加功能和未来前景
NeMo-RL v0.3引入了异步回滚和非同位置生成等功能，扩展了其能力。展望未来，NVIDIA计划支持更大的MOE模型，并引入进一步的优化，包括FP8生成支持和与Megatron-Core的非同位置生成。
NeMo-RL与Megatron-Core后端的进步标志着在优化大型语言模型的强化学习方面迈出了一大步，确保了模型训练的效率和可扩展性。
Image source: Shutterstock.jpg)