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. Los modelos MoE fragmentan sus expertos con "ep_dispatch_experts". Sus implementaciones determinan cómo llegan los tokens a los expertos. Token dispatch permite que ep_size difiera de tp_size. Algunas implementaciones de EP requieren ep_size=tp_size, por ejemplo "all_reduce" y "deepgemm_megamoe", de modo que cada rank de un grupo de expertos mantiene los mismos tokens.

Llama 4 no tiene un plan de EP: sus expertos procesan cada token ya escalado por su puntuación de enrutador, por lo que no hay nada que despachar. En su lugar, divídelos con 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 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, de modo que cada rank solo carga sus expertos locales, y la regla "ep_dispatch_experts" en el módulo de expertos selecciona cómo llegan los tokens a ellos.

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.

"ep_dispatch_experts" fragmenta expertos completos por rango (ep_size es un múltiplo de tp_size que divide a fsdp_size * tp_size) y accede a ellos a través de la implementación que indica ep_implementation, consulte Elegir una implementación.

Elegir una implementación

"ep_dispatch_experts" fragmenta los expertos una sola vez y determina cómo llegan los tokens a ellos a partir de las implementaciones registradas en EPImplementationInterface, seleccionadas con ep_implementation en DistributedConfig:

ep_implementation solo elige entre los tamaños de paralelismo definidos en DistributedConfig. Cuando los rangos de un grupo de expertos mantienen lotes diferentes, por ejemplo tp_size=1 con fsdp_size=ep_size, "all_to_all" es la única implementación aplicable.

Agregar una implementación

Una implementación subclase EPImplementation y se registra bajo el nombre ep_implementation selects. Se construye una vez por módulo de expertos, por lo que puede mantener estado, como buffers de comunicación. Al llamarla con el lote que posee este rank, con los ids globales de expertos, devuelve la salida ponderada de los expertos para cada uno de sus tokens. Ejecuta los expertos de este rank con experts_forward, sobre las filas, los ids locales de expertos y los pesos que les proporcionó.

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

Los gradientes deben llegar a cada rank que contenga un token, incluidos los pesos de enrutamiento, por lo que debe envolver una operación colectiva que autograd no diferencia en un torch.autograd.Function.

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.

Con ep_size=tp_size, la implementación predeterminada "auto" ejecuta "all_reduce" en la inferencia y distribuye los tokens en el entrenamiento. Para enmascarar y aplicar all-reduce también en el entrenamiento, establezca ep_implementation="all_reduce":

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

Despacho de tokens

El despacho de tokens es la implementación de "all_to_all": cada rank entrena con su propia parte del lote. En cada capa MoE, un rank enruta sus tokens, envía cada par (token, experto) al rank que posee el experto mediante un all-to-all, ejecuta sus expertos locales sobre lo que recibe, obtiene los resultados 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 rank 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 sus propias mallas, tp y ep, 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 otra dimensión de malla con fsdp_size. Con enmascarado y all-reduce, mantén ep_size=tp_size y pasa 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)

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