Resumo Técnico: Paralelismo Heterogêneo para Treinamento de Modelos de Linguagem Multimodais Grandes

1. Declaração do Problema

À medida que o treinamento de modelos fundamentais evolui de transformadores apenas de texto para sistemas multimodais (abrangendo imagens, áudio, vídeo e documentos), o grafo de treinamento torna-se cada vez mais heterogêneo em termos de profundidade, largura e tamanho de contexto dos módulos. No entanto, as pilhas de paralelismo atuais (Megatron-LM, DeepSpeed, etc.) tipicamente impõem um layout de paralelismo único e compartilhado (Paralelismo de Tensor, Contexto, Pipeline, Dados e Especialistas) em todo o grafo de ponta a ponta.

Essa abordagem de "tamanho único" cria ineficiências significativas quando os requisitos do codificador e do Modelo de Linguagem Grande (LLM) divergem:

• Comunicação Desnecessária: Grandes LLMs frequentemente exigem altos graus de Paralelismo de Tensor (TP). Forçar pequenos codificadores a herdar esses graus de TP introduz coletivas de all-reduce desnecessárias ao redor de computações locais do codificador, adicionando sobrecarga que supera as economias de memória.
• Incompatibilidade de Paralelismo de Contexto: O Paralelismo de Contexto (CP) é essencial para que LLMs gerenciem a memória de atenção sobre sequências longas. No entanto, os comprimentos de sequência do codificador são limitados pela resolução da modalidade (por exemplo, patches de imagem), e não pela janela de contexto do LLM. Herdar shards de CP do LLM para sequências de tamanho fixo do codificador não proporciona alívio de memória, enquanto esses ranks poderiam funcionar de forma mais eficiente como réplicas independentes de Paralelismo de Dados (DP) para processar mais amostras.
• Restrições de Pipeline: O Paralelismo de Pipeline (PP) frequentemente confina codificadores à primeira etapa do pipeline do LLM. Isso força o codificador a herdar o grau de TP do LLM e pode sobrecarregar ranks específicos, impedindo que o LLM utilize um layout de pipeline mais eficiente.

O desafio central é que o ponto de operação ótimo para computação, memória e comunicação muda com base nos parâmetros da carga de trabalho (escala do codificador, comprimento de contexto, densidade de tokens, tamanho do mundo). Um layout estático e compartilhado não consegue se adaptar a essas mudanças.

2. Metodologia

Os autores propõem Paralelismo Heterogêneo, uma abstração de tempo de execução que permite que módulos dentro de um único grafo multimodal de ponta a ponta escolham layouts lógicos independentes e posicionamentos físicos de ranks. O sistema suporta dois modos de execução distintos:

A. Abstração Central: Comunicadores de Fronteira

Para desacoplar layouts de módulos enquanto preserva a semântica de treinamento, o sistema introduz comunicadores de fronteira. Estes são operadores conscientes de autograd anexados às arestas entre módulos que lidam com a transformação de tensores entre diferentes representações distribuídas.

• Passagem Direta (Forward Pass): Materializa a ativação de origem no layout do módulo de destino.
• Passagem Reversa (Backward Pass): Aplica a transformação inversa para rotear gradientes de volta ao layout do módulo de origem.
• Corretude: O sistema garante que a semântica dos tensores seja preservada em grades independentes. Mapeamentos incorretos são prevenidos definindo estritamente transformações lógicas de shard de lote (Equal-DP, Fan-in e Fan-out) para cada aresta.

B. Modos de Execução

1. Execução Colocalizada:

   • Mecanismo: Módulos compartilham as mesmas GPUs físicas, mas as interpretam sob grades lógicas diferentes (por exemplo, um codificador com alto DP e baixo TP compartilhando ranks com um LLM com baixo DP e alto TP).
   • Implementação: O comunicador atua como um operador de layout de rank compartilhado. Ele executa operações coletivas locais (por exemplo, all-gather ou seleção de intervalo) para reinterpretar o layout de dados sem mover tensores entre conjuntos de ranks.
   • Agendamento: Utiliza um cronograma de três fases para evitar conflitos entre coletivas do codificador e o pipeline escalonado do LLM (1F1B).
     • Fase 1: Forward do codificador + transformação forward colocalizada.
     • Fase 2: Execução de pipeline do LLM desacoplada (1F1B) no tensor transformado.
     • Fase 3: Transformação backward colocalizada + backward do codificador.
   • Benefício: Permite que parâmetros do codificador sejam descarregados durante a fase do pipeline do LLM, melhorando a eficiência de memória.

2. Execução Não Colocalizada:

   • Mecanismo: Módulos ocupam conjuntos disjuntos de ranks (por exemplo, uma "ilha de codificador" e uma "ilha de LLM").
   • Implementação: Utiliza comunicação explícita Ponto a Ponto (P2P) entre líderes de fronteira determinísticos (ranks com coordenadas canônicas).
     • Forward: Líderes de origem enviam para líderes de destino; transmissões locais materializam os dados.
     • Backward: Gradientes são roteados de volta através dos caminhos inversos dos líderes.
   • Agendamento: Utiliza um dispatcher 1F1B consciente de grafo que resolve arestas em um grafo multietapa, lidando com junções e ramificações onde múltiplos fluxos de codificador alimentam o LLM.
   • Benefício: Previne que o estado do codificador (parâmetros, ativações, estados do otimizador) compita com o LLM por memória nos mesmos ranks, permitindo que o LLM utilize um layout de pipeline mais eficiente.

3. Principais Contribuições

1. Abstração de Runtime de Paralelismo Heterogêneo: Um sistema que permite layouts independentes de TP/CP/PP/DP/EP e posicionamentos físicos de ranks para módulos em um único grafo de MLLM, suportando tanto execução colocalizada quanto não colocalizada.
2. Comunicadores de Fronteira e Agendamento: A introdução de comunicadores de fronteira conscientes de autograd que lidam com transformações de layout, acoplados a agendadores especializados (três fases para colocalizado, 1F1B consciente de grafo para não colocalizado) para manter a eficiência do pipeline.
3. Avaliação de Regime de Operação: Uma avaliação abrangente caracterizando quando a heterogeneidade é benéfica. O estudo identifica que a heterogeneidade colocalizada é ótima quando há folga de memória suficiente, mas os layouts lógicos diferem, enquanto a heterogeneidade não colocalizada torna-se superior quando o estado do codificador restringe a memória do LLM ou o layout do pipeline.
4. Corretude e Código Aberto: Validação da paridade de convergência de perda contra baselines homogêneas por meio de verificações FP32 em nível de etapa e varreduras de treinamento estilo LLaVA. O sistema é lançado como uma extensão de código aberto para o Megatron-LM.

4. Resultados Experimentais

O sistema foi avaliado em GPUs NVIDIA H100 de 80GB em várias cargas de trabalho multimodais (codificadores de 1B–12B, LLMs de 14B–120B) e comprimentos de contexto (8K–128K).

• Heterogeneidade Colocalizada:
  • Melhora o TFLOPS/GPU em até 49,3% (observado em cenários de contexto longo, por exemplo, 64K–128K, onde LLMs exigem CP, mas codificadores não).
  • Melhora o throughput em 12,1%–21,6% em várias configurações de codificador 1B/3B + LLM 14B/70B/120B, permitindo que o codificador use DP mais alto e TP mais baixo.
• Heterogeneidade Não Colocalizada:
  • Melhora o throughput agregado de tokens em até 13,0% e o TFLOPS/GPU em até 9,6% para cargas de trabalho de codificador grande (6B–12B) + LLM 120B.
  • Torna-se o ponto de operação preferido quando a escala do codificador (por exemplo, 6B+) cria pressão de memória que restringe o layout do pipeline do LLM em configurações colocalizadas.
• Convergência: Validado que layouts heterogêneos produzem curvas de perda idênticas às baselines homogêneas, confirmando que os comunicadores de fronteira preservam corretamente a semântica dos gradientes.

5. Significado e Alegações

O artigo afirma que o paralelismo heterogêneo é uma evolução necessária para o treinamento multimodal, afastando-se de restrições rígidas de layout em todo o modelo. Ao expor graus de liberdade adicionais, o sistema permite que a infraestrutura de treinamento ajuste o ponto de operação dinamicamente com base nas características da carga de trabalho (tamanho do codificador, comprimento de contexto, densidade de tokens).

Os autores enfatizam que isso não é meramente uma mudança de configuração, mas um substrato de tempo de execução fundamental que preserva a semântica de treinamento de ponta a ponta, permitindo:

• Execução colocalizada para extrair o máximo de throughput de um orçamento fixo de GPUs quando módulos podem compartilhar hardware, mas precisam de grades lógicas diferentes.
• Execução não colocalizada para desacoplar pressões de memória e agendamento quando o estado do codificador, de outra forma, criaria um gargalo para o LLM.

O trabalho é apresentado como uma extensão de código aberto ao Megatron-LM, fornecendo uma base reutilizável para futuras estruturas de MLLM modulares e sistemas de treinamento desagregados. O artigo não alega resolver todos os desafios multimodais, mas aborda especificamente a ineficiência causada por forçar módulos heterogêneos em um único layout de paralelismo.