Paralelismo de expertos

Paralelismo de expertos es una estrategia de paralelismo para modelos de mezcla de expertos (MoE). La capa de avance de cada experto reside en un acelerador de hardware diferente. Un enrutador envía los tokens a los expertos correspondientes y recopila los resultados. Este enfoque permite escalar los modelos a recuentos de parámetros mucho mayores sin aumentar el costo de cómputo, ya que cada token activa solo unos pocos expertos.

DistributedConfig

Active el paralelismo de expertos con la clase DistributedConfig y el argumento ep_size. La mayoría de los modelos MoE usan token dispatch de forma predeterminada, por lo que ep_size se puede establecer de forma independiente de tp_size. Algunos, como Llama 4 y Gemma 4, siguen usando enmascarado y all-reduce de forma predeterminada ("ep_router" y "moe_tp_experts" en model.ep_plan). El enmascarado también está disponible en cualquier modelo con una sobrescritura de ep_plan y requiere ep_size=tp_size para que cada rank de un grupo de expertos reciba los mismos 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 dos planes en su configuración: base_model_tp_plan para los módulos densos y base_model_ep_plan para los expertos. Se exponen en el modelo cargado como model.tp_plan y model.ep_plan. Con tp_size > 1 y ep_size > 1, se aplican ambos: el plan de paralelismo de tensor divide la atención y los MLP densos, y el plan de paralelismo de expertos divide los expertos. Las reglas de EP tienen prioridad sobre las reglas de TP para los mismos módulos, por lo que los pesos de los expertos se dividen una sola vez, mediante el plan de EP. En el plan de EP, el estilo GroupedGemmParallel divide los pesos de los expertos a lo largo de la dimensión de expertos para que cada rank cargue solo sus expertos locales, y ep_router enmascara los expertos que residen en otros ranks antes de que un all-reduce combine las salidas de los expertos.

tp_plan se aplica solo cuando tp_size > 1, y ep_plan solo cuando ep_size > 1. Con TP habilitado y EP deshabilitado, se aplica el plan de TP completo, incluidas las reglas de expertos.

La regla de avance de expertos en ep_plan selecciona cómo llegan los tokens a los expertos:

enable_expert_parallel=True sin ep_size está en desuso y emite una FutureWarning. Solo fragmenta los expertos, como ep_size=tp_size, tp_plan={}.

Lance su script de inferencia con torchrun. El número de procesos debe ser igual a tp_size * fsdp_size * pp_size, y ep_size debe dividir exactamente el número de expertos.

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

Sobrescribir los planes

Pase tp_plan={...} o ep_plan={...} a DistributedConfig para sobrescribir reglas individuales de los planes predefinidos. Las reglas no especificadas se conservan y los planes fusionados se almacenan en el modelo. Cada clave debe coincidir con un módulo, un parámetro o una entrada existente del plan; de lo contrario, la carga genera un ValueError antes de que se fragmente nada. Use la ruta completa tal como se ve desde el modelo cargado, por lo tanto model.layers.* para un LM causal y layers.* para su modelo base.

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"},
)

Proporcionar un plan no infiere los tamaños de paralelismo: establezca tp_size y ep_size explícitamente.

La mayoría de los modelos MoE usan "ep_dispatch_experts" por defecto. Para usar enmascarado y all-reduce en su lugar, establezca ep_size=tp_size y anule tanto las reglas de forward del router como las de los expertos (el nombre del módulo del router depende del modelo, p. ej. mlp.router en 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 el contrario, sobrescribe la regla de forward de los expertos de un modelo cuyo plan utiliza enmascarado con "ep_dispatch_experts" para usar el despacho de tokens. La regla del router se ignora en ese caso, ya que el despacho necesita los ids globales de los expertos para encontrar el propietario de cada experto.

Despacho de tokens

Con el despacho de tokens, cada rango entrena con su propia parte del lote. En cada capa MoE, un rango enruta sus tokens, envía cada par (token, experto) al rango que posee el experto mediante un all-to-all, ejecuta sus expertos locales sobre lo que recibe, obtiene los resultados de vuelta con un segundo all-to-all y los combina con los pesos de enrutamiento. Solo viajan las activaciones enrutadas y las salidas de los expertos, y ningún rango calcula expertos para tokens que no posee.

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)

Con tp_size=1, ep_size debe dividir a fsdp_size y al número de expertos, y la atención no está limitada por num_key_value_heads. Para el resto del modelo:

• Los parámetros que están fuera de los expertos se fragmentan con FSDP2 a través de fsdp, lo que reduce sus gradientes.
• Los expertos se fragmentan en ep y, cuando efsdp_size = fsdp_size * tp_size // ep_size es mayor que uno, además se fragmentan con FSDP en efsdp. Siempre están envueltos con FSDP, por lo que fsdp_mixed_precision y fsdp_cpu_offload también se aplican a ellos y save_pretrained() los reúne como cualquier otro parámetro.
• Un grupo de paralelismo de expertos contiene ep_size / tp_size lotes, por lo que el gradiente de un experto es la suma sobre esa cantidad de lotes. La reducción de efsdp divide por fsdp_size en lugar de por el tamaño de su grupo, lo que produce el mismo promedio por lote que FSDP2 utiliza para los módulos densos.
• Cada experto local también procesa una fila de relleno de ceros por capa. Un rango cuyos expertos no recibieron tokens aún participa en el reverse all-to-all y en la reducción de gradientes de expertos.

Cómo se combinan los tamaños

• tp_size * fsdp_size es el número de procesos. ep_size no añade ninguno: agrupa los mismos rangos solo para los pesos de los expertos.
• ep_size divide la lista de expertos en ep_size bloques. Cada rank calcula num_experts / ep_size expertos, y ep_size ranks consecutivos mantienen un conjunto completo. Ese conjunto de ranks es el grupo en el que se ejecuta el all-to-all.
• efsdp_size = fsdp_size * tp_size / ep_size es el número de copias completas del conjunto de expertos. Los rangos en la misma posición en copias diferentes fragmentan esos expertos para la memoria y promedian sus gradientes, como FSDP lo hace para los módulos densos.
• El lote que mantiene un rango depende únicamente de fsdp. Los rangos consecutivos forman un grupo de TP y reciben el mismo lote; los grupos fsdp_size reciben lotes diferentes.

Ocho procesos, tp_size=2, fsdp_size=4, ocho expertos:

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

De la imagen se desprenden dos números:

• Dentro de un grupo de EP, cada token existe tp_size veces, una por cada rango del par que mantiene su lote. Esto no depende de ep_size.
• Un grupo EP mantiene ep_size / tp_size lotes diferentes. Esta es la cantidad sobre la que se suma el gradiente de un experto.

Con paralelismo de tensores

Establezca tp_size > 1 para fragmentar los módulos densos con el plan de TP mientras que los expertos usan dispatch. En ocho procesos:

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

Cada par de rangos de TP recibe el mismo lote, porque la paralelización de tensores replica las activaciones dentro del par. Si ambos rangos despacharan todos sus tokens, el rango propietario de un experto recibiría cada token dos veces, lo calcularía dos veces y su gradiente de pesos se duplicaría. Los expertos están completos en un solo rango, por lo que el duplicado no puede dividirse por pesos, y el propietario suele ser otro rango, por lo que no puede resolverse por propiedad como lo hace la enmascaración. Por lo tanto, el par divide las filas: cada rango de TP despacha una fracción disjunta de 1 / tp_size de los tokens, los resultados vuelven al rango que los envió y un all-reduce sobre el par de las mitades rellenas con ceros restaura la salida replicada que la siguiente capa espera. La división es por tp_size, no por ep_size, ya que solo los rangos que poseen un lote pueden enviarlo. Los grupos de expertos abarcan cuatro rangos y cada experto está fragmentado con FSDP en efsdp_size = 2 rangos, mientras que el grupo de FSDP del tronco abarca cuatro rangos. Las restricciones habituales de TP del modelo, como la divisibilidad de las cabezas de atención, siguen aplicándose a los módulos densos. Las rebanadas de tokens pueden ser desiguales o vacías, incluso durante la decodificación de un solo token.

El despacho de tokens aún no se puede combinar con la paralelización de pipeline; use pp_size=1 (aún no probado) Estas configuraciones utilizan ocho GPUs cada una:

El Trainer aún no tiene en cuenta el despacho de tokens: el conteo de lotes y de tokens asume el diseño de all-reduce, donde los rangos de un grupo de TP comparten un lote y existen fsdp_size fragmentos de paralelismo de datos. El soporte del Trainer para el despacho llegará en una actualización posterior.

Combinación con FSDP2

El paralelismo de tensores y de expertos fragmenta los pesos en tp, pero el estado del optimizador y los módulos sin una regla siguen replicados en cada rango del grupo, lo que limita el tamaño del modelo que se puede entrenar. Añade FSDP2 en una segunda dimensión de malla con fsdp_size. Con enmascarado y all-reduce, mantén ep_size=tp_size y pasa `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)

Cargue el modelo como de costumbre y luego entrene con Trainer. Toma la norma del gradiente en ambas mallas y asigna a cada malla su propio grupo de parámetros del optimizador. save_model() reúne los pesos fragmentados en un punto de control regular. Esto requiere accelerate>=1.12 para que Trainer pueda reflejar tp_size y fsdp_size en ~Accelerate.ParallelismConfig.

La tabla siguiente compara el entrenamiento solo con EP con el entrenamiento 2D EP+FSDP2 en 8xH100 GPUs. La carga de trabajo es el ajuste completo de Qwen3-30B-A3B en bf16 con una longitud de secuencia de 2048. Más fragmentos de FSDP reducen la memoria pico, y los tokens/s disminuyen un poco porque FSDP2 realiza all-gathers y reduce-scatters de los expertos a través de fsdp.

Reanudar desde un punto de control aún no se admite para modelos fragmentados en el momento de la carga, por lo que el Trainer solo acepta save_only_model=True o save_strategy="no" para ellos.

Referencia de la API