Expert Parallelism

O paralelismo de especialistas (EP) é uma estratégia de treino distribuído e inferência para modelos de mistura de especialistas (MoE), que coloca diferentes especialistas em diferentes GPUs e encaminha cada token para os dispositivos que contêm os especialistas selecionados para o processar. Uma camada MoE substitui um grande bloco feed-forward de um transformador por várias redes de especialistas mais pequenas e por um encaminhador que ativa apenas alguns especialistas por token, pelo que o total de parâmetros pode crescer muito mais depressa do que a computação aplicada a cada token. O paralelismo de especialistas permite armazenar esses parâmetros na memória: em vez de cada GPU armazenar todos os especialistas, cada uma armazena uma parte. Já em 2017, Noam Shazeer e colegas da Google distribuíram especialistas pelos dispositivos, com uma cópia partilhada de cada um, na sua camada MoE esparsamente controlada; em junho de 2020, o GShard ampliou esta abordagem, distribuindo os especialistas de um modelo de tradução com mais de 600 mil milhões de parâmetros por 2,048 aceleradores TPU v3. O paralelismo de especialistas é hoje padrão no treino e na execução de modelos de linguagem e modelos de visão-linguagem MoE de grande dimensão.

Como os tokens circulam entre GPUs#

Cada camada MoE é executada em quatro etapas:

1. Encaminhamento: O encaminhador pontua cada token em relação a cada especialista e seleciona os especialistas top-k, normalmente um, dois ou oito.
2. Encaminhamento: Uma operação coletiva de todos para todos envia o estado oculto de cada token para as GPUs que contêm os especialistas selecionados.
3. Computação: Cada GPU executa os seus especialistas locais nos tokens recebidos.
4. Combinação: Uma segunda operação all-to-all devolve os resultados à GPU original de cada token, onde são ponderados pelas pontuações do encaminhador e somados.

O excerto abaixo simula as etapas de encaminhamento e distribuição de uma camada com 8 especialistas distribuídos por 4 GPUs, dois especialistas por GPU:

import torch

torch.manual_seed(0)
num_experts, num_gpus, top_k = 8, 4, 2
tokens = torch.randn(4096, 512)  # estados ocultos de 4.096 tokens
router = torch.nn.Linear(512, num_experts)

# Encaminhamento: cada token escolhe os seus 2 melhores especialistas com base na pontuação do encaminhador
with torch.no_grad():
    _, experts = router(tokens).topk(top_k, dim=-1)

# Paralelismo de especialistas: os especialistas 0-1 estão na GPU 0, os especialistas 2-3 na GPU 1 e assim por diante
gpu = experts // (num_experts // num_gpus)

# Distribuição: contar as atribuições de tokens a especialistas recebidas por cada especialista e GPU
print(torch.bincount(experts.flatten(), minlength=num_experts).tolist())
print(torch.bincount(gpu.flatten(), minlength=num_gpus).tolist())

A primeira linha do resultado mostra a carga de cada especialista, entre 948 e 1.087 atribuições, face a uma quota igual de 1.024. A segunda linha mostra que as quatro GPUs recebem entre 1.968 e 2.139 das 8.192 atribuições. Mesmo um encaminhador inicializado aleatoriamente distribui o trabalho de forma desigual, e os encaminhadores treinados desenvolvem especialistas populares que desequilibram ainda mais a carga.

Equilíbrio de carga e capacidade#

Uma camada só termina quando a GPU mais ocupada termina; por isso, um encaminhamento desequilibrado deixa o hardware inativo. O treino MoE combate este problema de três formas:

• Perda auxiliar de equilíbrio: Um termo adicional na função de perda recompensa o encaminhador por distribuir os tokens uniformemente pelos especialistas.
• Fator de capacidade: Cada especialista aceita, no máximo, um número fixo de tokens por lote — por exemplo, 1.25 vezes a quota equitativa — e os tokens excedentes ignoram esse especialista.
• Equilíbrio baseado em desvios: O DeepSeek-V3 acrescenta um desvio por especialista às pontuações do encaminhador, que afeta apenas a seleção dos especialistas. Após cada passo, reduz o desvio dos especialistas sobrecarregados e aumenta o dos que estão inativos, equilibrando a carga sem uma perda auxiliar que concorra com o objetivo principal.

Com um limite de capacidade, o percurso de cada token depende da carga dos seus especialistas:

Um token ignorado atravessa a ligação residual da camada sem alterações, pelo que a eliminação de muitos tokens prejudica a qualidade. O DeepSeek-V3 não elimina tokens durante o treino e, em vez disso, recorre ao equilíbrio baseado em desvios.

A comunicação é o outro custo. As duas trocas de todos para todos ocorrem em cada camada MoE; por isso, o paralelismo de especialistas funciona melhor dentro de um nó, onde as GPUs partilham ligações NVLink rápidas, e torna-se dispendioso através das ligações mais lentas entre nós. O DeepSeek-V3 distribui os seus 256 especialistas encaminhados por camada com paralelismo de especialistas de 64 vias em 8 nós, envia cada token para especialistas em, no máximo, 4 nós e sobrepõe as trocas à computação.

Combinação de estratégias de paralelismo#

O paralelismo de especialistas é uma das várias formas de distribuir o treino por vários dispositivos:

• Paralelismo de dados: Replica o modelo completo em cada GPU e divide cada lote, sincronizando os gradientes após cada passo.
• Paralelismo de tensores: Divide matrizes de pesos individuais entre GPUs, para que todas as GPUs participem em cada camada.
• Paralelismo de pipeline: Atribui camadas consecutivas a GPUs diferentes e processa microlotes em fluxo através dessas camadas.
• Paralelismo de contexto: Divide uma única sequência de entrada longa entre GPUs.

Os modelos densos que cabem em uma GPU não precisam de nada além de paralelismo de dados, que é como Ultralytics YOLO26 treina em várias GPUs quando uma lista de dispositivos, como device=[0, 1], é passada no modo de treinamento. Modelos MoE grandes combinam várias formas de paralelismo: o DeepSeek-V3 treinou seus 671 bilhões de parâmetros, com 37 bilhões ativos por token, usando paralelismo de pipeline de 16 vias, paralelismo de especialistas de 64 vias e paralelismo de dados fragmentado, sem paralelismo de tensores. Mecanismos de inferência como o vLLM aplicam a mesma divisão ao servir um modelo, permitindo que um modelo MoE cujos especialistas excederiam a memória de uma GPU seja executado em várias GPUs, cada uma contendo apenas sua parte dos especialistas.