Treinamento de contexto longo do MoE

Documentação estável: @docs/training/moe-optimization.md Cartão: @skills/nemo-mbridge-perf-moe-long-context/card.yaml

O que muda em contextos longos

Quando o comprimento da sequência ultrapassa amplamente a faixa de 4K, a memória de atenção e a permanência da ativação tornam-se as restrições dominantes. Para modelos MoE, isso geralmente significa que é necessária alguma combinação de:

• paralelismo de contexto
• recálculo seletivo
• menor precisão
• descarregamento da CPU para o estado do otimizador
• um despachante e um layout PP que não desperdicem o orçamento de DP restante, que é menor

Padrões de escalonamento arredondados

DSV3 no H100

As execuções de contexto longo do DSV3 apresentam um padrão estável:

• o recálculo seletivo funciona melhor do que o recálculo completo assim que se ultrapassam os contextos mais curtos
• a taxa de transferência permanece em uma faixa bastante estreita, desde contextos de comprimento médio até os muito longos, se o CP for aumentado adequadamente
• o equilíbrio muda de “capacidade de memória” para “viabilidade do número de GPUs” à medida que o CP aumenta

Em outras palavras, o contexto longo não reduz imediatamente a utilização se o layout for bem escolhido, mas consome o orçamento de DP muito rapidamente.

Qwen3-Next no GB200

O Qwen3-Next se comporta mais como um modelo de média escala sensível à memória:

• 8K e 32K continuam sendo viáveis com CP moderado
• 64K é possível, mas a queda na taxa de transferência é perceptível e a memória fica muito mais limitada
• melhorias no layout do pipeline e no GEMM agrupado são quase tão importantes quanto o CP

Qwen3 235B no GB200

O Qwen3 235B mostra que um contexto longo ainda pode ser eficiente em sistemas NVL72 quando TP, CP e HybridEP são coordenados. As melhores configurações da classe 128K não são apenas receitas do tipo “apenas adequadas”; elas podem permanecer altamente eficientes se o roteamento, o paralelismo e o recálculo estiverem equilibrados.

Regras práticas para dimensionamento de CP

1. Comece com uma meta de fragmento de 4K: uma boa estimativa inicial é CP ~= seq_len / 4096, depois arredonde para um layout prático que seja uma potência de dois.

2. Mantenha o DP ativo, se possível: o escalonamento de contexto longo se torna instável quando CP, EP, TP e PP, juntos, reduzem o DP ao mínimo.

3. Prefira a recomputação seletiva: recompute módulos como up_proj, norm, moe, moe_act ou mlp antes de recorrer à recomputação completa.

4. Evite recálculos com uso intenso de SDPA em contextos muito longos: recalcular os componentes internos da atenção pode adicionar muito trabalho para um benefício de memória menor do que recalcular módulos menores do lado do MoE e do MLP.

5. Use o TP como outra alavanca em sistemas NVL72: execuções no GB200 e no GB300 podem, às vezes, trocar um pouco de CP por TP sem perder a eficiência.

6. Presuma que o GBS precisará ser reduzido: à medida que o CP aumenta e o DP diminui, talvez seja necessário reduzir o tamanho do lote global ou aceitar um GA mais alto.

Famílias de configurações representativas

DSV3 a 128K no H100

TP=1  CP=32  EP=32  PP=8  VPP=4
Precisão: classe FP8
Despachador: DeepEP
Recálculo: up_proj, norm, moe, mlp
Ajuda extra de memória: descarregamento da CPU pelo otimizador

DSV3 com 256K no H100

TP=1  CP=64  EP=32  PP=8  EDP=2  VPP=4
Precisão: classe FP8
Despachador: DeepEP
Recálculo: up_proj, norm, moe, mlp
Ajuda extra de memória: descarregamento da CPU pelo otimizador

Qwen3 235B a 128K no GB200

TP=4  CP=4  EP=32  PP=4  VPP=12
Precisão: BF16 ou MXFP8
Despachante: HybridEP
Recálculo: moe_act, norm
Grafo CUDA: attn + moe_router + moe_preprocess

Orientação sobre recalculação e gráfico CUDA

Para treinamento MoE de contexto longo:

• comece com o recálculo seletivo
• adicione gráficos CUDA somente depois que as formas e o caminho de roteamento estiverem estáveis
• mantenha o comprimento da sequência e o MBS fixos ao usar gráficos CUDA
• se a execução depender de lotes altamente dinâmicos, dê preferência à execução antecipada

Referências úteis:

• @docs/training/activation-recomputation.md
• @skills/nemo-mbridge-perf-cuda-graphs/SKILL.md

Armadilhas

1. O CP não substitui o EP nem o PP: ele acrescenta outra dimensão; não faz com que os outros desapareçam.

2. Uma boa linha de base 4K ainda pode ser uma linha de base ruim para contextos longos: o modo de roteamento, a escolha de recálculo e a estratégia de descarregamento muitas vezes precisam ser alterados.

3. A viabilidade em termos de número de GPUs torna-se a verdadeira restrição: um contexto muito longo pode parecer adequado em uma única receita, mas tornar-se impossível quando o EP e o PP são adicionados de forma consistente em todo o modelo.

4. Os gráficos CUDA precisam de formatos estáticos: lotes de comprimento variável e estratégias oportunistas de preenchimento podem interromper o caminho sem aviso prévio.

5. O suporte a contêineres e kernels é mais importante em 128K+: caminhos de contexto longo tendem a depender de kernels mais recentes e correções de bugs do que a inicialização de contexto curto.