Paralelismo de especialistas

Paralelismo de especialistas é uma estratégia de paralelismo para modelos mixture-of-experts (MoE). A camada feedforward de cada especialista reside em um acelerador de hardware diferente. Um roteador despacha os tokens para os especialistas apropriados e reúne os resultados. Essa abordagem escala os modelos para contagens de parâmetros muito maiores sem aumentar o custo de computação, pois cada token ativa apenas alguns especialistas.

DistributedConfig

Ative o paralelismo de especialistas com a classe DistributedConfig e o argumento ep_size. A maioria dos modelos MoE usa token dispatch por padrão, portanto ep_size pode ser definido independentemente de tp_size. Alguns, como Llama 4 e Gemma 4, ainda usam masking e all-reduce por padrão ("ep_router" e "moe_tp_experts" em model.ep_plan). O masking também está disponível em qualquer modelo com um override de ep_plan e exige ep_size=tp_size para que cada rank em um grupo de especialistas receba os mesmos tokens.

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    distributed_config=distributed_config,
)

Cada modelo MoE define dois planos em sua configuração: base_model_tp_plan para os módulos densos e base_model_ep_plan para os especialistas. Eles são expostos no modelo carregado como model.tp_plan e model.ep_plan. Com tp_size > 1 e ep_size > 1, ambos se aplicam: o plano de paralelismo de tensor faz o sharding da atenção e dos MLPs densos, e o plano de paralelismo de especialistas faz o sharding dos especialistas. As regras de EP têm precedência sobre as regras de TP para os mesmos módulos, de modo que os pesos dos especialistas são shardados uma única vez, pelo plano de EP. No plano de EP, o estilo GroupedGemmParallel divide os pesos dos especialistas ao longo da dimensão de especialista, de modo que cada rank carrega apenas seus especialistas locais, e o ep_router mascara os especialistas que estão em outros ranks antes de um all-reduce combinar as saídas dos especialistas.

tp_plan é aplicado apenas quando tp_size > 1, e ep_plan apenas quando ep_size > 1. Com o TP habilitado e o EP desabilitado, o plano TP completo é aplicado, incluindo as regras de especialista.

A regra de encaminhamento de especialistas em ep_plan seleciona como os tokens chegam aos especialistas:

enable_expert_parallel=True sem ep_size está obsoleto e emite um FutureWarning. Ele faz o sharding apenas dos experts, como em ep_size=tp_size, tp_plan={}.

Inicie seu script de inferência com torchrun. O número de processos deve ser igual a tp_size * fsdp_size * pp_size, e ep_size deve dividir uniformemente o número de especialistas.

torchrun --nproc-per-node 8 your_script.py

Sobrescrevendo os planos

Passe tp_plan={...} ou ep_plan={...} para DistributedConfig para sobrescrever regras individuais dos planos predefinidos. As regras não especificadas são mantidas e os planos mesclados são armazenados no modelo. Cada chave deve corresponder a um módulo, a um parâmetro ou a uma entrada existente do plano; caso contrário, o carregamento levanta um ValueError antes que qualquer coisa seja fragmentada. Use o caminho completo conforme visto a partir do modelo carregado, ou seja, model.layers.* para um LM causal e layers.* para o modelo base correspondente.

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    tp_plan={"model.layers.*.self_attn.q_proj": "colwise_rep"},
    ep_plan={"model.layers.*.mlp.experts.down_proj": "grouped_gemm"},
)

Fornecer um plano não infere os tamanhos de paralelismo: defina tp_size e ep_size explicitamente.

A maioria dos modelos MoE usa "ep_dispatch_experts" por padrão. Para usar mascaramento e all-reduce em vez disso, defina ep_size=tp_size e substitua tanto as regras de forward do roteador quanto as dos especialistas (o nome do módulo do roteador depende do modelo, por exemplo, mlp.router no gpt-oss):

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)

Por outro lado, substitua a regra de forward dos especialistas de um modelo cujo plano usa mascaramento com "ep_dispatch_experts" para usar o despacho de tokens. A regra do roteador é então ignorada, pois o despacho precisa dos ids globais dos especialistas para encontrar o proprietário de cada especialista.

Despacho de tokens

Com o token dispatch, cada rank treina em sua própria parte do lote. Em cada camada MoE, um rank roteia seus tokens, envia cada par (token, especialista) para o rank que possui o especialista por meio de um all-to-all, executa seus especialistas locais no que recebe, obtém os resultados de volta com um segundo all-to-all e os combina com os pesos de roteamento. Apenas as ativações roteadas e as saídas dos especialistas trafegam, e nenhum rank calcula especialistas para tokens que não possui.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=1,
    fsdp_size=8,
    ep_size=4,
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

Com tp_size=1, ep_size deve dividir fsdp_size e o número de especialistas, e a atenção não é limitada por num_key_value_heads. Para o restante do modelo:

• Os parâmetros fora dos especialistas são fragmentados com FSDP2 em fsdp, o que reduz seus gradientes.
• Os especialistas são fragmentados em ep e, quando efsdp_size = fsdp_size * tp_size // ep_size é maior que um, também são fragmentados via FSDP em efsdp. Eles sempre são envolvidos por FSDP, portanto fsdp_mixed_precision e fsdp_cpu_offload também se aplicam a eles e save_pretrained() os agrega como qualquer outro parâmetro.
• Um grupo de paralelismo de especialistas contém ep_size / tp_size lotes, portanto, o gradiente de um especialista é a soma sobre esse número de lotes. A redução efsdp divide por fsdp_size em vez do tamanho do seu grupo, o que resulta na mesma média por lote que o FSDP2 utiliza para os módulos densos.
• Cada especialista local também processa uma linha de zero pad por camada. Um rank cujos especialistas não receberam tokens ainda participa do reverse all-to-all e da redução de gradiente dos especialistas.

Como os tamanhos se combinam

• tp_size * fsdp_size é o número de processos. ep_size não adiciona nenhum: ele apenas reagrupa os mesmos ranks para os pesos dos especialistas.
• ep_size divide a lista de especialistas em ep_size blocos. Cada rank calcula num_experts / ep_size especialistas, e ep_size ranks consecutivos mantêm um conjunto completo. Esse conjunto de ranks é o grupo em que o all-to-all é executado.
• efsdp_size = fsdp_size * tp_size / ep_size é o número de cópias completas do conjunto de especialistas. Os ranks na mesma posição em cópias diferentes fazem o sharding desses especialistas para memória e fazem a média de seus gradientes, assim como o FSDP faz para os módulos densos.
• O lote que um rank mantém depende apenas de fsdp. Ranks consecutivos formam um grupo de TP e recebem o mesmo lote; os grupos fsdp_size recebem lotes diferentes.

Oito processos, tp_size=2, fsdp_size=4, oito especialistas:

rank         0      1       2      3       4      5       6      7
batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]     one batch per TP pair
tp           0      1       0      1       0      1       0      1

ep_size=2   E0-3   E4-7    E0-3   E4-7    E0-3   E4-7    E0-3   E4-7     group = a TP pair,   efsdp_size=4
ep_size=4   E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7     group = two pairs,   efsdp_size=2
ep_size=8   E0     E1      E2     E3      E4     E5      E6     E7       group = all ranks,   efsdp_size=1

Dois números decorrem da imagem:

• Dentro de um grupo de EP, cada token existe tp_size vezes, uma vez por rank do par que detém seu lote. Isso não depende de ep_size.
• Um grupo EP mantém ep_size / tp_size lotes diferentes. Essa é a quantidade sobre a qual a gradiente de um especialista é somada.

Com paralelismo de tensor

Defina tp_size > 1 para dividir os módulos densos com o plano de TP, enquanto os especialistas usam dispatch. Em oito processos:

distributed_config = DistributedConfig(
    tp_size=2,
    fsdp_size=4,
    ep_size=4,
)

Cada par de ranks de TP recebe o mesmo batch, pois o paralelismo de tensor replica as ativações dentro do par. Se ambos os ranks despachassem todos os seus tokens, o rank que possui um expert receberia cada token duas vezes, o calcularia duas vezes e o gradiente do seu peso seria duplicado. Como os experts estão inteiros em um único rank, a duplicata não pode ser dividida pelos pesos, e como o proprietário geralmente é outro rank, o problema não pode ser resolvido pela propriedade, como a máscara faz. O par, portanto, divide as linhas: cada rank de TP despacha uma fração 1 / tp_size disjunta dos tokens, os resultados retornam ao rank que os enviou e um all-reduce sobre o par das metades com preenchimento zero restaura a saída replicada que a próxima camada espera. A divisão é por tp_size, não por ep_size, pois apenas os ranks que possuem um batch podem enviá-lo. Os grupos de experts abrangem quatro ranks e cada expert é fragmentado por FSDP em efsdp_size = 2 ranks, enquanto o grupo FSDP do tronco abrange quatro ranks. As restrições de TP habituais do modelo, como a divisibilidade das cabeças de atenção, ainda se aplicam aos módulos densos. As fatias de tokens podem ser desiguais ou vazias, inclusive durante a decodificação de um único token.

O despacho de tokens ainda não pode ser combinado com paralelismo de pipeline; use pp_size=1 (ainda não testado) Cada uma dessas configurações usa oito GPUs:

O Trainer ainda não considera o despacho de tokens: a contagem de lotes e de tokens pressupõe o layout de all-reduce, em que os ranks de um grupo de TP compartilham um lote e existem fsdp_size shards de paralelismo de dados. O suporte do Trainer para despacho será implementado em uma atualização posterior.

Combinando com FSDP2

O paralelismo de tensor e de especialistas divide os pesos entre tp, mas o estado do otimizador e os módulos sem uma regra ainda são replicados em cada rank do grupo, o que limita o tamanho do modelo que você pode treinar. Adicione FSDP2 em uma segunda dimensão de malha com fsdp_size. Com mascaramento e all-reduce, mantenha ep_size=tp_size e passe `ep_plan={“layers..mlp.gate”: “ep_router”, “layers..mlp.experts”: “moe_tp_experts”}“.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,  # expert parallel size, must match tp_size with masking and all-reduce
    fsdp_size=2,  # data parallel shards
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

Carregue o modelo como de costume e, em seguida, treine com Trainer. Ele calcula a norma do gradiente em ambas as malhas e atribui a cada malha seu próprio grupo de parâmetros do otimizador. save_model() reúne os pesos fragmentados em um checkpoint regular. Isso requer accelerate>=1.12 para que o Trainer possa espelhar tp_size e fsdp_size em ~Accelerate.ParallelismConfig.

A tabela abaixo compara o treinamento apenas com EP e o 2D EP+FSDP2 em 8 GPUs H100. A carga de trabalho é o fine-tuning completo do Qwen3-30B-A3B em bf16, com comprimento de sequência 2048. Mais shards do FSDP reduzem a memória de pico, e os tokens/s caem um pouco porque o FSDP2 faz all-gathers e reduce-scatters dos experts através do fsdp.

Retomar a partir de um checkpoint ainda não é suportado para modelos fragmentados no momento do carregamento, portanto, o Trainer aceita apenas save_only_model=True ou save_strategy="no" para eles.

Referência da API