Xiaohongshu apresenta o BigMac: superando o conflito entre memória e velocidade no treinamento multimodal

Os modelos de linguagem multimodais de grande porte estão se tornando a pedra angular da IA de próxima geração; no entanto, sua infraestrutura de treinamento tem sido há muito tempo prejudicada por um dilema persistente: os sistemas otimizados para velocidade frequentemente enfrentam restrições de memória, enquanto aqueles que priorizam a eficiência de memória tendem a ser lentos. A equipe Dots Infra da Xiaohongshu identifica esse gargalo como a fronteira de Pareto do treinamento em pipeline multimodal e agora conseguiu superá-lo. Em 22 de julho, a equipe disponibilizou como código aberto o BigMac, um novo paradigma de treinamento paralelo em pipeline projetado especificamente para cenários multimodais nativos. O projeto já está disponível no GitHub no repositório Dots-Infra.

Ao contrário dos transformadores padrão, os modelos multimodais são inerentemente complexos. Um Modelo de Linguagem Grande Multimodal (MLLM) típico compreende três componentes distintos: codificadores modais que convertem imagens e áudio em representações, uma estrutura central de LLM para inferência e um gerador que mapeia as saídas do LLM de volta para as modalidades-alvo, como imagens ou fala. As diferenças estruturais significativas entre esses componentes criam desafios substanciais ao integrá-los em um único pipeline de treinamento.

As soluções atuais do setor geralmente se enquadram em duas categorias. A primeira abordagem prioriza a eficiência computacional ao desacoplar codificadores e geradores do pipeline do LLM, executando-os separadamente. Isso evita que flutuações na duração dos módulos modais criem “bolhas” no pipeline do LLM, mas faz com que a memória de ativação aumente proporcionalmente ao número de micro-lotes, levando a altos custos em escala. A segunda abordagem se concentra na eficiência de memória, mantendo todos os módulos dentro do mesmo pipeline, o que encurta os ciclos de vida das ativações e reduz o uso de memória. No entanto, se qualquer codificador ou gerador for lento, todo o pipeline de LLM fica paralisado, resultando em “bolhas de cauda”. À medida que a escala do modelo aumenta, ambos os projetos revelam gargalos críticos.

O BigMac aborda esse desafio com um princípio simples, mas fundamental: o pipeline central do LLM continua sendo o foco principal. O treinamento de LLMs em grande escala já conta com estratégias maduras de agendamento, como 1F1B ou 1F1B intercalado, que estão profundamente integradas em pilhas de treinamento de nível de produção. Em vez de substituir esses métodos estabelecidos, o BigMac usa o cronograma do LLM como sua linha do tempo subjacente, inserindo estrategicamente os cálculos do codificador e do gerador quando as entradas estão prontas e sem interromper a ordem de execução do LLM. A equipe se refere a essa arquitetura como um “pipeline aninhado quase seguro contra dependências”.

Esse projeto oferece duas vantagens principais. Primeiro, as flutuações na duração do codificador e do gerador não se propagam mais pelo pipeline do LLM, permitindo que o LLM mantenha seu ritmo ideal. Segundo, os requisitos de memória para ativações modais são reduzidos a O(1) no nível algorítmico. Os codificadores não precisam mais reter ativações para todos os micro-lotes até que o pipeline seja concluído, e os geradores evitam prolongar longas caudas de ativação. Essencialmente, o BigMac não troca memória por picos de desempenho; em vez disso, ele reestrutura a lógica de agendamento para tornar esses dois objetivos compatíveis, em vez de mutuamente exclusivos.

Preencher a lacuna entre o projeto do agendamento e a execução real envolve desafios de engenharia significativos, incluindo integração de sistemas, definição de interfaces e depuração de desempenho. O BigMac foi projetado para lidar com esses desafios específicos, oferecendo três recursos principais. Primeiro, ele torna o agendamento global transparente: o agendador gera uma tabela abrangente de operadores que cobre todas as instâncias do pipeline, micro-lotes e tipos de módulos em tempo de execução. O executor, então, distribui esses elementos em sequências locais para cada rank, integrando-se perfeitamente a back-ends de LLM como o Megatron Core, ambientes de execução modais e camadas de comunicação. Essa visibilidade facilita a verificação e a otimização do back-end.

Em segundo lugar, o BigMac oferece uma interface de paralelismo de pipeline que permanece invisível para os engenheiros de algoritmos. Os usuários precisam apenas definir o que cada módulo produz e consome; o sistema lida internamente com a divisão de estágios, ativação e transferência de gradientes, além da comunicação entre dispositivos. Isso permite que experimentos multimodais verificados em uma única GPU sejam escalonados naturalmente para o paralelismo de pipeline. Em terceiro lugar, ele fornece um conjunto de ferramentas de perfilagem, simulação e visualização. Essas ferramentas detalham as iterações de treinamento até o nível do operador, revelando exatamente o que cada rank está fazendo a cada passo temporal, identificando períodos de inatividade e destacando gargalos de dependência. O simulador também permite que as equipes testem várias configurações de PP e combinações de micro-lotes antes de iniciar o treinamento em escala real, estimando seu impacto nas bolhas e na taxa de processamento.

A eficácia do BigMac foi validada em duas cargas de trabalho representativas. Em tarefas de compreensão de MLLM, utilizando o Qwen3-30B-A3B como estrutura base e um codificador ViT de 1,3B, o BigMac alcança um aumento de velocidade de 1,08x a 1,1x em comparação com o Optimus, que é a linha de base computacionalmente eficiente, e um aumento de 1,6x a 1,9x em relação ao Megatron-DistTrain, que é a linha de base eficiente em termos de memória. Fundamentalmente, o uso de memória permanece estável à medida que o tamanho do lote por GPU aumenta, enquanto o Optimus sofre picos de consumo de memória que acabam levando a erros de falta de memória (OOM) em lotes maiores. As tarefas de geração MLLM são mais complexas, envolvendo um gerador MMDiT de 20 bilhões, onde a diferença de desempenho se amplia ainda mais: o Optimus falha em todos os tamanhos de lote testados devido a erros de OOM, enquanto o BigMac evita isso ao executar o gerador de forma eficiente em sentido reverso e liberar rapidamente as ativações. Em comparação com o Megatron-DistTrain, o BigMac ainda oferece um aumento de velocidade de 1,5x a 1,9x com uso estável de memória. Isso destaca o principal valor dos pipelines aninhados: atender às dependências tanto do codificador quanto do gerador sem retenção excessiva de ativações ou tempo de inatividade significativo.

O BigMac é agora um componente central da infraestrutura de treinamento de modelos multimodais da Dots e opera no ambiente de produção do Xiaohongshu. O artigo que acompanha o projeto, “BigMac: Breaking the Pareto Frontier of Compute and Memory in Multimodal LLM Training”, está disponível no arXiv, juntamente com exemplos interativos de rastreamento do PP Profiler. Para pesquisadores e engenheiros que enfrentam o dilema entre memória e velocidade no treinamento multimodal, este projeto de código aberto validado em produção oferece uma solução prática que economiza tempo e esforço significativos.