Paralelismo de modelo é uma técnica de computação distribuída usada em aprendizado de máquina e aprendizado profundo para treinar ou executar inferência em redes neurais grandes que não cabem na memória de um único dispositivo, como uma GPU ou TPU. Diferentemente do paralelismo de dados, em que cada dispositivo mantém uma cópia completa do modelo e processa diferentes lotes de dados, o paralelismo de modelo particiona o próprio modelo, tipicamente atribuindo diferentes camadas ou subconjuntos de camadas a diferentes dispositivos. Essa abordagem é essencial para escalar modelos como os transformers usados em modelos de linguagem de grande porte, que podem ter bilhões ou trilhões de parâmetros.
Visão geral
No paralelismo de modelo, a rede neural é dividida ao longo de sua profundidade (por camadas) ou, em formas mais avançadas, ao longo de outras dimensões, como largura ou até mesmo operações individuais de tensores. A forma mais direta é o particionamento por camadas, em que camadas consecutivas são colocadas em dispositivos diferentes. Durante a propagação direta, as ativações fluem de um dispositivo para o próximo; durante a retropropagação, os gradientes fluem na direção inversa. Isso cria uma execução semelhante a um pipeline, que pode introduzir tempo ocioso (bolhas) se não for cuidadosamente agendada. O paralelismo de modelo é frequentemente contrastado com o paralelismo de dados, mas os dois podem ser combinados em abordagens híbridas, como o paralelismo 3D usado no treinamento de modelos grandes.
Motivação
As redes neurais modernas, particularmente os modelos de linguagem de grande porte, cresceram exponencialmente em tamanho. Por exemplo, o GPT-3 (2020) tem 175 bilhões de parâmetros, exigindo mais de 350 GB de memória apenas para os pesos em precisão de 32 bits, excedendo em muito a capacidade de 80 GB de GPUs de ponta como a A100. Mesmo com otimizações de memória como precisão mista e checkpointing de gradientes, o treinamento em um único dispositivo é inviável. O paralelismo de modelo permite que pesquisadores treinem modelos ordens de magnitude maiores do que um único dispositivo pode armazenar, distribuindo o modelo por um cluster de dispositivos conectados por interconexões de alta velocidade.
Tipos de paralelismo de modelo
Paralelismo por camadas (pipeline)
Esta é a forma mais simples, em que o modelo é dividido em estágios sequenciais, cada um atribuído a um dispositivo diferente. Por exemplo, uma rede de 100 camadas poderia ser dividida em 10 estágios de 10 camadas cada, com cada estágio em uma GPU separada. Durante a passagem direta, os dados fluem pelo estágio 1, depois pelo estágio 2, e assim por diante. O principal desafio é o balanceamento de carga e a redução de bolhas no pipeline. Técnicas como GPipe e PipeDream introduzem micro-lotes e agendamento para melhorar a utilização.
Paralelismo de tensores
O paralelismo de tensores divide operações individuais (por exemplo, multiplicações de matrizes) entre vários dispositivos. Por exemplo, no mecanismo de atenção de um transformer, as projeções de query, key e value podem ser particionadas entre GPUs, com os resultados combinados por meio de operações all-reduce. Essa abordagem é usada no Megatron-LM e é eficaz para reduzir a memória por dispositivo, mantendo alta eficiência computacional, mas requer comunicação de alta largura de banda.
Paralelismo de especialistas
Usado em modelos de mistura de especialistas (MoE), o paralelismo de especialistas coloca diferentes redes especialistas em dispositivos diferentes, enquanto o roteador (rede de portas) distribui tokens para os especialistas apropriados. Essa é uma forma de paralelismo de modelo que permite contagens massivas de parâmetros sem aumento proporcional de computação, pois apenas um subconjunto de especialistas é ativado por token.
Desafios de implementação
O paralelismo de modelo introduz vários desafios:
- Sobrecarga de comunicação: Os dispositivos precisam trocar ativações e gradientes, o que pode se tornar um gargalo, especialmente com interconexões lentas.
- Desequilíbrio de carga: Tamanhos de camadas ou requisitos computacionais desiguais podem fazer com que alguns dispositivos fiquem ociosos enquanto outros trabalham.
- Bolhas no pipeline: No paralelismo de pipeline, os dispositivos podem esperar por dados de estágios anteriores, reduzindo a utilização.
- Fragmentação de memória: O particionamento de modelos pode levar a uso desigual de memória, exigindo posicionamento e agendamento cuidadosos.
- Tolerância a falhas: Com muitos dispositivos, a probabilidade de falha aumenta, exigindo mecanismos de checkpointing e recuperação.
Comparação com paralelismo de dados
No paralelismo de dados, cada dispositivo mantém uma cópia completa do modelo e processa um mini-lote diferente de dados. Os gradientes são calculados pela média entre os dispositivos após cada etapa. Isso é simples e escala bem para modelos que cabem na memória de um único dispositivo. O paralelismo de modelo, por outro lado, é necessário quando o próprio modelo é grande demais. No entanto, o paralelismo de modelo frequentemente tem custos de comunicação mais altos e menor eficiência devido a dependências seriais. Na prática, o treinamento em larga escala usa ambos: paralelismo de dados entre nós e paralelismo de modelo dentro de um nó.
Aplicações
O paralelismo de modelo é fundamental para treinar modelos de linguagem de grande porte como GPT-4, PaLM e LLaMA, bem como para servi-los em produção. Empresas como OpenAI, Google (via Google Cloud) e Anthropic dependem do paralelismo de modelo para treinar e implantar seus modelos. Também é usado em outros domínios, como visão computacional (por exemplo, CNNs 3D) e computação científica, onde os modelos são grandes demais para um único acelerador.
Suporte de software
Vários frameworks fornecem suporte integrado para paralelismo de modelo:
- PyTorch: Oferece
torch.distributedcom paralelismo de pipeline (por exemplo,torch.distributed.pipeline.sync.Pipe) e paralelismo de tensores por meio de bibliotecastensor_parallel. - TensorFlow: Fornece
tf.distributecom estratégias comoTPUStrategyeMultiWorkerMirroredStrategyque podem incorporar paralelismo de modelo. - Megatron-LM: Uma biblioteca especializada da NVIDIA para paralelismo de tensores e pipeline.
- DeepSpeed: Da Microsoft, oferece ZeRO (Zero Redundancy Optimizer), que é uma forma de paralelismo de modelo que particiona estados do otimizador, gradientes e parâmetros.
- JAX: Com
jax.sharding, permite posicionamento explícito de dispositivos e execução paralela.
Considerações de hardware
O paralelismo de modelo depende fortemente de comunicação de alta velocidade entre dispositivos. Dentro de um único nó, NVLink e PCIe fornecem conexões rápidas; entre nós, InfiniBand ou Ethernet de alta velocidade são usados. Hardware especializado como GPUs NVIDIA e TPUs Google frequentemente inclui interconexões dedicadas (por exemplo, NVSwitch, links de TPU) para facilitar o paralelismo de modelo. Empresas como Cerebras e SambaNova projetaram sistemas com grande memória no chip para reduzir a necessidade de paralelismo de modelo, mas para escala extrema, abordagens distribuídas permanecem necessárias.
Avanços recentes
Pesquisas recentes têm se concentrado em melhorar a eficiência do paralelismo de modelo. Técnicas como paralelismo de sequência (dividindo a dimensão de sequência) e paralelismo de contexto (para sequências longas) foram desenvolvidas. Algoritmos automáticos de particionamento de modelos, como os do Alpa e FlexFlow, otimizam o posicionamento de operações entre dispositivos. Além disso, o paralelismo híbrido combinando dados, tensores e pipeline é agora padrão no treinamento em larga escala, como exemplificado pelo modelo Megatron-Turing NLG (530B parâmetros) treinado pela NVIDIA e Microsoft.
Limitações e direções futuras
Apesar de seu poder, o paralelismo de modelo tem limitações. A sobrecarga de comunicação pode dominar para modelos pequenos ou redes lentas. A complexidade do particionamento e agendamento aumenta com o tamanho do modelo e a heterogeneidade do cluster. Direções futuras incluem auto-paralelismo mais inteligente, melhor compressão de comunicação e novas arquiteturas que são inerentemente mais paralelizáveis. Com o avanço da IA generativa e modelos cada vez maiores, o paralelismo de modelo permanecerá uma pedra angular do aprendizado profundo escalável.
Ver também
- Paralelismo de dados
- Paralelismo de pipeline
- Paralelismo de tensores
- Mistura de especialistas
- Treinamento distribuído
- Modelo de linguagem de grande porte
- Cluster de GPUs
Referências
- Narayanan, D., et al. (2019). PipeDream: Generalized Pipeline Parallelism for DNN Training.
- Shoeybi, M., et al. (2019). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.
- Rajbhandari, S., et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models.
- Huang, Y., et al. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.