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. Modelos MoE dividem seus especialistas com "ep_dispatch_experts". Suas implementações definem como os tokens chegam aos especialistas. O Token dispatch permite que ep_size difira de tp_size. Algumas implementações de EP exigem ep_size=tp_size, por exemplo, "all_reduce" e "deepgemm_megamoe", para que cada rank de um grupo de especialistas mantenha os mesmos tokens.

Llama 4 não possui um plano de EP: seus especialistas recebem todos os tokens já escalados pela pontuação do roteador, portanto, não há nada para despachar. Em vez disso, divida-os com paralelismo de tensor (tp_size).

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, para que cada rank carregue apenas seus especialistas locais, e a regra "ep_dispatch_experts" no módulo de especialistas seleciona como os tokens chegam até eles.

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.

"ep_dispatch_experts" faz o sharding de especialistas inteiros por rank (ep_size é um múltiplo de tp_size que divide fsdp_size * tp_size) e os alcança por meio da implementação nomeada por ep_implementation, veja Escolhendo uma implementação.

Escolhendo uma implementação

"ep_dispatch_experts" faz o sharding dos experts uma única vez e define como os tokens os alcançam a partir das implementações registradas em EPImplementationInterface, selecionadas com ep_implementation em DistributedConfig:

ep_implementation apenas escolhe entre os tamanhos de paralelismo definidos em DistributedConfig. Quando os ranks de um grupo de especialistas mantêm lotes diferentes, por exemplo tp_size=1 com fsdp_size=ep_size, "all_to_all" é a única implementação aplicável.

Adicionando uma implementação

Uma implementação é uma subclasse de EPImplementation e é registrada sob o nome selecionado por ep_implementation. Ela é construída uma vez por módulo de especialistas, o que permite que ela mantenha estado, como buffers de comunicação. Ao chamá-la com o lote que este rank possui, com os ids globais dos especialistas, ela retorna a saída ponderada dos especialistas para cada um de seus tokens. Ela executa os especialistas deste rank com experts_forward, nas linhas, ids locais dos especialistas e pesos que ela trouxe.

from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface

class MyEP(EPImplementation):
    def __call__(self, hidden_states, top_k_index, top_k_weights):
        tokens, expert_ids, weights = ...  # bring this rank's experts the rows routed to them
        expert_output = self.experts_forward(tokens, expert_ids, weights)
        return ...  # the weighted expert output for each token of the batch this rank holds

EPImplementationInterface.register("my_ep", MyEP)
distributed_config = DistributedConfig(tp_size=1, fsdp_size=8, ep_size=8, ep_implementation="my_ep")

Os gradientes devem alcançar todos os ranks que mantêm um token, incluindo os pesos de roteamento, portanto, envolva uma operação coletiva que o autograd não diferencia em um torch.autograd.Function.

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.

Com ep_size=tp_size, a implementação padrão "auto" executa "all_reduce" na inferência e distribui tokens no treinamento. Para mascarar e fazer all-reduce também no treinamento, defina ep_implementation="all_reduce":

distributed_config = DistributedConfig(tp_size=4, ep_size=4, ep_implementation="all_reduce")

Despacho de tokens

O despacho de tokens é a implementação "all_to_all": 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 em suas próprias malhas, tp e ep, 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 outra dimensão da malha com fsdp_size. Com mascaramento e all-reduce, mantenha ep_size=tp_size e passe ep_implementation="all_reduce".

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_implementation="all_reduce",  # "auto" would dispatch tokens in training, which the Trainer does not support yet
)
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