Atraso de gradiente de uma etapa não é uma barreira para o pré-treinamento de LLM em pipeline paralelo assíncrono em larga escala

Resumo

O pré-treinamento moderno de LLMs em larga escala se beneficia do uso de Pipeline Parallelism; no entanto, implementações síncronas deixam GPUs ociosas durante as bolhas do pipeline, desperdiçando recursos computacionais. O Pipeline Parallelism Assíncrono elimina essas bolhas, maximizando a vazão ao custo de desatualização do gradiente. Entre os esquemas assíncronos, o PipeDream-2BW é particularmente atraente: diferentemente do esquema PipeDream original, ele garante um atraso de gradiente constante de uma etapa, independentemente da profundidade do pipeline. No entanto, sua adoção permanece limitada devido à crença comum de que a otimização sob desatualização é fundamentalmente instável. Neste trabalho, desafiamos essa suposição, demonstrando que a degradação sob atraso de uma etapa depende fortemente da escolha do otimizador, em vez de ser uma limitação intrínseca. Fornecemos a primeira análise empírica abrangente mostrando que, embora o AdamW, o otimizador predominante na época em que o PipeDream-2BW foi introduzido, de fato sofra degradação severa, métodos recentes como o Muon exibem robustez forte sob um atraso de uma etapa. Introduzimos uma correção inspirada em Error Feedback, agnóstica ao otimizador, para mitigar ainda mais os efeitos do atraso. Oferecemos análise teórica de suporte demonstrando convergência para o Muon com e sem essa correção. Uma avaliação extensa em modelos de até 10 bilhões de parâmetros confirma que nossas estratégias reduzem a lacuna de desempenho com o treinamento síncrono, destacando o potencial prático do pipeline parallelism assíncrono em larga escala.

English

Modern large-scale LLM pretraining benefits from utilizing Pipeline Parallelism; however, synchronous implementations leave GPUs idle during pipeline bubbles, wasting computational resources. Asynchronous Pipeline Parallelism eliminates these bubbles, maximizing throughput at the cost of gradient staleness. Among asynchronous schedules, PipeDream-2BW is particularly appealing: unlike the original PipeDream schedule, it ensures a constant one-step gradient delay regardless of pipeline depth. However, its adoption remains limited due to the common belief that optimizing under staleness is fundamentally unstable. In this work, we challenge this assumption, demonstrating that degradation under one-step delay depends strongly on optimizer choice rather than being an intrinsic limitation. We provide the first comprehensive empirical analysis showing that while AdamW, the predominant optimizer at the time when PipeDream-2BW was introduced, indeed suffers from severe degradation, recent methods like Muon exhibit strong robustness under a one-step delay. We introduce an optimizer-agnostic Error Feedback-inspired correction to further mitigate delay effects. We provide supporting theoretical analysis demonstrating convergence for Muon with and without this correction. Extensive evaluation on models up to 10B parameters confirms that our strategies bridge the performance gap with synchronous training, highlighting the practical potential of asynchronous pipeline parallelism at scale.