Pipeline parallelism é uma estratégia de treinamento distribuído para redes neurais profundas que particiona as camadas de um modelo em estágios sequenciais, cada um atribuído a um dispositivo diferente ( como GPUs ou TPUs).. Ela combina paralelismo de modelo com paralelismo de dados ao processar micro-lotes através dos estágios de forma semelhante a um pipeline de software, visando aumentar a taxa de transferência de treinamento e reduzir o tempo ocioso (bolhas) que o paralelismo ingênuo por camadas causaria
Contexto e motivação
Treinar grandes redes neurais, especialmente grandes modelos de linguagem e outros modelos de aprendizado profundo, requer recursos computacionais enormes. Conforme os tamanhos dos modelos crescem para bilhões ou trilhões de parâmetros, o treinamento em um único dispositivo torna-se inviável devido a limites de memória e computação. Paralelismo de dados replica o modelo inteiro em cada dispositivo e divide os dados, mas falha quando o modelo excede a memória de um único dispositivo. Paralelismo de modelo divide o próprio modelo entre dispositivos, mas a colocação ingênua por camadas leva a subutilização severa: apenas um dispositivo computa por vez enquanto os outros esperam, resultando em baixa taxa de transferência e alto tempo ocioso
Pipeline parallelism aborda isso dividindo o modelo em estágios ( tipicamente grupos contíguos de camadas)e processando múltiplos micro-lotes concorrentemente. A ideia-chave é sobrepor a computação entre estágios: enquanto o estágio 1 computa no micro-lote 2, o estágio 2 pode computar no micro-lote 1, e assim por diante. Isso se assemelha a uma linha de montagem, onde cada estágio processa continuamente dados de entrada e passa os resultados para o próximo estágio
Conceitos centrais
A unidade fundamental de trabalho em pipeline parallelism é o micro-lote. Um lote de treinamento grande é dividido em micro-lotes menores que fluem através dos estágios do pipeline sequencialmente. Cada estágio realiza passagens para frente e para trás em suas camadas atribuídas. O pipeline é "preenchido" alimentando micro-lotes um após o outro, e após um período inicial de aquecimento, todos os estágios podem estar ocupados simultaneamente, alcançando alta taxa de transferência
A eficiência de um pipeline é frequentemente medida pela razão de bolha, que é a fração de tempo em que os estágios estão ociosos. Para um pipeline com p estágios e m micro-lotes, a taxa de transferência ideal é alcançada quando m é muito maior que p, minimizando a sobrecarga relativa de bolhas. No entanto, aumentar m também aumenta o uso de memória para ativações, criando um trade-off
Principais esquemas
Vários esquemas de pipeline parallelism foram propostos, diferindo em como agendam passagens para frente e para trás e como lidam com atualizações de gradientes.
GPipe
Google introduziu GPipe em 2019. Ele particiona o modelo em estágios e usa um agendamento simples: cada micro-lote realiza passagens para frente através de todos os estágios, depois passagens para trás através de todos os estágios em ordem reversa. GPipe usa atualizações de gradientes síncronas, significando que todos os estágios esperam todos os micro-lotes completarem antes de atualizar os pesos. Isso garante gradientes consistentes mas introduz uma bolha que cresce com o número de estágios. GPipe também requer ou re-computação de ativações ou armazenamento delas, trocando computação por memória
PipeDream
PipeDream, da Microsoft Research (2019), usa um agendamento assíncrono onde cada estágio processa micro-lotes de forma round-robin, realizando passagens para frente e para trás o mais rápido possível. Isso reduz a bolha mas introduz desatualização de pesos: diferentes estágios podem usar versões diferentes dos pesos, o que pode prejudicar a convergência. PipeDream também usa uma técnica chamada "weight stashing" para manter múltiplas versões de pesos para diferentes micro-lotes, aumentando a sobrecarga de memória
PipeDream-2BW e variantes
PipeDream-2BW (2020) melhora o PipeDream ao usar dois buffers de pesos (um para frente, um para trás) para reduzir memória e melhorar convergência. Variantes como V-Pipe e PipeMare exploram diferentes estratégias de sincronização
Megatron-LM e agendamento intercalado
NVIDIA's Megatron-LM (2019) combinou paralelismo de tensor (dividindo camadas individuais) com pipeline parallelism. Posteriormente, Megatron-2 (2020) introduziu um agendamento intercalado que divide o modelo em mais estágios do que dispositivos, permitindo que cada dispositivo lide com múltiplos estágios. Isso reduz o tamanho da bolha ao aumentar o número de micro-lotes que podem estar em voo, ao custo de mais comunicação
Chimera e outros esquemas recentes
Chimera (2021) usa um pipeline bidirecional para reduzir ainda mais bolhas. Outros esquemas como PTD-P (2021) e ZeroBubble (2023) exploram agendamentos mais sofisticados para eliminar ou reduzir bolhas de pipeline
Considerações matemáticas e práticas
Pipeline parallelism é frequentemente combinado com outras estratégias de paralelismo. Paralelismo de tensor divide camadas individuais entre dispositivos, enquanto pipeline parallelism divide o modelo verticalmente. Paralelismo de dados replica o pipeline entre grupos de dispositivos. Essa abordagem híbrida é padrão no treinamento de grandes modelos, como visto em sistemas como Megatron-Turing NLG e OpenAI's GPT-4.
A escolha da profundidade do pipeline (número de estágios)e do tamanho do micro-lote afeta o desempenho. Pipelines mais profundos reduzem a memória por estágio mas aumentam comunicação e sobrecarga de bolhas. Micro-lotes maiores melhoram a utilização mas consomem mais memória. A configuração ótima depende do tamanho do modelo, da memória do dispositivo, e da largura de banda da interconexão
Aplicações e impacto
Pipeline parallelism tem sido instrumental no treinamento de alguns dos maiores modelos de IA. Por exemplo, GPT-3 (175B parâmetros) foi treinado usando uma combinação de paralelismo de modelo e pipeline parallelism. Anthropic's Claude modelos e Google's PaLM também dependem de técnicas semelhantes. A abordagem é suportada por frameworks principais, incluindo PyTorch (através do módulo torch.distributed.pipeline), TensorFlow (via tf.distribute), e JAX (com pjit)
Desafios e limitações
Pipeline parallelism introduz vários desafios. A sobrecarga de comunicação surge do envio de ativações e gradientes entre estágios, o que pode se tornar um gargalo em interconexões lentas. O desequilíbrio de carga ocorre quando estágios têm tempos de computação desiguais, reduzindo a eficiência. A pressão de memória do armazenamento de ativações para passagens para trás pode ser mitigada por re-computação de ativações, mas a um custo computacional. Além disso, esquemas assíncronos podem sofrer de problemas de convergência devido a pesos desatualizados, enquanto esquemas síncronos incorrem em bolhas
Direções futuras
A pesquisa continua em melhorar pipeline parallelism, focando em reduzir bolhas, minimizar memória, e adaptar a hardware heterogêneo. Técnicas como particionamento automático de estágios, balanceamento dinâmico de carga, e integração com modelos de mistura de especialistas são áreas ativas. O surgimento de modelos extremamente grandes e a necessidade de treinamento eficiente em clusters de nuvem provavelmente manterão pipeline parallelism um componente-chave dos sistemas de treinamento distribuído.