NVIDIA NeMo-RL Megatron-Core: Aumento de Desempenho de Treinamento

Visão Geral

O projeto NVIDIA NeMo-RL (NeMo Reinforcement Learning) inicializou com suporte a treinamento via PyTorch DTensor (conhecido como FSDP2), permitindo integração nativa com HuggingFace e escalonamento com paralelismos nativos do PyTorch como FSDP2, paralelismo de tensor, paralelismo de sequência e paralelismo de contexto. À medida que os modelos atingem centenas de bilhões de parâmetros, o caminho DTensor se torna inadequado devido ao consumo de memória de ativação e ao overhead de recomputação, o que leva a tempos de passo proibitivamente lentos, além da ausência de kernels CUDA otimizados e de outras melhorias de desempenho. Isso mostrou a necessidade de um backend mais eficiente, que o Megatron-Core oferece. Na versão NeMo-RL v0.3, a NVIDIA adiciona suporte ao backend Megatron-Core, com documentação detalhada, scripts de exemplo e arquivos de configuração para pós-treinamento de grandes modelos de forma eficiente. Megatron-Core traz técnicas otimizadas pela GPU e aprimoramentos de alto desempenho, estruturados em uma estratégia de paralelismo de 6D que otimiza padrões de comunicação e computação, além de suportar uma grande variedade de arquiteturas de modelo. O NeMo-RL agora permite que desenvolvedores usem essas otimizações durante o pós-treinamento. O NeMo-RL simplifica o uso do Megatron-Core ao automatizar grande parte da complexa sintonia por trás das cenas, expondo uma superfície de configuração mais simples e intuitiva. Ativar o treinamento baseado em Megatron é direto: adicione a seção policy.megatron_cfg no seu arquivo YAML e defina enabled: true. Um exemplo completo de configuração está disponível na documentação; todos os argumentos dentro do config serão encaminhados ao Megatron durante o treinamento. Depois de adicionar a seção megatron e definir enabled=True, você está pronto para treinar um modelo. Iniciar o treinamento é feito da mesma forma que com o DTensor, conforme descrito no README ou no guia para reproduzir o DeepScaleR. O treinamento baseado em Megatron-Core oferece suporte a modelos densos e Mixture-of-Experts (MoE). Os resultados incluem quebras de tempo de passo para GRPO (Group Relative Policy Optimization) em vários modelos, demonstrando desempenho superior ao DTensor, mantendo as mesmas propriedades de convergência. Observa-se o uso de kernels otimizados pela GPU, paralelismo 4D, packing de sequências e amostragem de importância como fatores de ganho de taxa de transferência, com treinamento de contexto longo suportado até 16K tokens. O paralelismo de contexto com Megatron-Core e DTensor também é suportado para treinamentos de contexto longo, incluindo resultados para modelos grandes como o Llama 3.3 70B. A conclusão é que o treinamento com Megatron-Core melhora significativamente a taxa de transferência de RL para modelos muito grandes, mantendo as propriedades de convergência. Para pesquisadores e engenheiros que buscam pós-treinamento de RL em modelos muito grandes, o Megatron-Core oferece uma rota para maior taxa de transferência, suporte a contexto longo e cobertura de modelos mais ampla, apoiada por documentação, configurações de exemplo e scripts para facilitar a adoção. A documentação do NeMo RL, configurações de exemplo e scripts ajudam a iniciar o pós-treinamento de seus grandes modelos com as otimizações Megatron-Core.

Principais características

Suporte ao backend Megatron-Core no NeMo-RL v0.3, permitindo otimizações no pós-treinamento para modelos grandes.
kernels otimizados pela GPU e melhorias de alto rendimento.
Estratégia de paralelismo 6D para otimizar comunicação e computação; suporta uma ampla variedade de arquiteturas de modelo.
Paralelismo 4D em conjunto com Megatron-Core, além de recursos como packing de sequências e amostragem de importância.
Suporte a treinamento de contexto longo, com testes até 16K sequências.
Paralelismo de contexto com Megatron-Core e DTensor para treinamentos de contexto longo.
Suporte a modelos densos e MoE (Mixture-of-Experts).
Configuração simplificada via policy.megatron_cfg; sintonia automática de grande parte das opções por trás das cenas.
Compatível com fluxos de trabalho DTensor existentes; a inicialização de treinamento segue o caminho DTensor.
Documentação, exemplos de configuração completos e scripts para ajudar a reproduzir resultados como os mostrados para o Llama 3.3 70B.

Casos de uso Comuns

Pós-treinamento de RL para modelos de linguagem extremamente grandes (com tamanhos na faixa de centenas de bilhões de parâmetros).
Cenários que exigem treinamento de contexto longo (testado até 16K comprimento de sequência) para manter contexto ao longo de fluxos de entrada longos.
Experimentos e fluxos de trabalho de produção envolvendo modelos densos e também arquiteturas MoE.
Experimentação rápida e integração com o ecossistema HuggingFace, aproveitando as otimizações do Megatron-Core para taxa de transferência.

Instalação & Configuração

Para habilitar o Megatron-Core no NeMo-RL, adicione a configuração Megatron ao seu YAML:

policy.megatron_cfg:
enabled: true

Isso segue a orientação de que o treinamento baseado em Megatron é ativado adicionando a seção megatron no config e definindo enabled: true. Todos os argumentos dentro do config serão encaminhados ao Megatron durante o treinamento. Após adicionar a seção megatron e habilitá-la, você pode treinar. Iniciar o treinamento é feito da mesma forma que com o DTensor, conforme descrito no README ou no guia de reprodução DeepScaleR. O treinamento baseado em Megatron-Core também suporta modelos densos e MoE. Para um início mínimo, garanta que a seção Megatron esteja presente e habilitada no config, e siga o fluxo de treinamento DTensor documentado no README do projeto.

Quick start

Garanta que o Megatron esteja habilitado na configuração do NeMo-RL:

policy.megatron_cfg:
enabled: true

Inicie o treinamento usando o mesmo caminho de inicialização de DTensor. Consulte o guia DTensor e o README do NeMo-RL para os comandos exatos e os scripts para reproduzir o DeepScaleR.
Para um exemplo mínimo de início, utilize o snippet YAML acima e, a partir dele, consulte a documentação para o exemplo completo de configuração e os arquivos de script fornecidos.
Quanto à escolha de modelo, Megatron-Core suporta modelos densos e MoE; você pode aplicar as rotinas de treinamento com o backend Megatron habilitado.

Quick start (exemplo mínimo executável)

YAML mínimo para habilitar Megatron Core:

policy.megatron_cfg:
enabled: true

Inicie o treinamento usando as mesmas entradas DTensor descritas no README do NeMo-RL; os argumentos de configuração Megatron serão encaminhados automaticamente durante o treinamento.

Prós e Contras

Prós
Taxa de transferência de treinamento substancialmente maior para modelos grandes em comparação com o caminho PyTorch DTensor.
Kernels otimizados pela GPU e paralelismo 4D/6D para reduzir gargalos de comunicação e melhorar a eficiência de computação.
Suporte a treinamento de contexto longo (testado até 16K) e modelos de grande escala como Llama 3.3 70B.
Suporte a modelos densos e MoE.
Configuração simplificada por meio de policy.megatron_cfg; sintonia automática reduz a complexidade para novos usuários.
Paralelismo de contexto com Megatron-Core e DTensor facilita treinamentos de contexto longo.
Contras
Megatron-Core introduz muitas opções de configuração de baixo nível; embora o NeMo-RL automatize grande parte, usuários avançados ainda podem ajustar opções para workloads específicos.
A configuração envolve adicionar e habilitar configurações Megatron específicas no YAML, o que é uma etapa adicional além do fluxo DTensor básico.

Alternativas (comparações rápidas)

| Aspecto | PyTorch DTensor (FSDP2) | Megatron-Core (NeMo-RL) |---|---|---| | Throughput em modelos muito grandes | Passos mais lentos devido à memória de ativação e overhead de recomputação; kernels CUDA otimizados ausentes | Throughput mais alto com kernels otimizados da GPU e estratégia de paralelismo 6D |Propriedades de convergência | Mantém as propriedades de convergência padrão | Mantém as propriedades de convergência, otimizado para throughput |Complexidade / sintonia | Configuração de superfície mais simples, mas pode exigir sintonia manual para escala | Automatiza grande parte da sintonia; configuração mais simples via policy.megatron_cfg |Suporte a modelos | Suporte amplo com paralelismos PyTorch nativos | Suporte a modelos densos e MoE; contexto longo suportado |Comprimento de contexto | Dependente das capacidades do DTensor | Demonstração de treinamento de contexto longo até 16K; melhorias em contextos longos |