Parallélisme d'experts

Le parallélisme d’experts est une stratégie de parallélisation destinée aux modèles mixture-of-experts (MoE). La couche feedforward de chaque expert réside sur un accélérateur matériel différent. Un routeur achemine les jetons vers les experts appropriés et collecte les résultats. Cette approche permet de faire évoluer les modèles vers des nombres de paramètres bien plus élevés sans augmenter le coût de calcul, car chaque jeton n’active que quelques experts.

DistributedConfig

Activez le parallélisme d'experts avec la classe DistributedConfig et l'argument ep_size. La plupart des modèles MoE utilisent par défaut le token dispatch, ce qui permet de définir ep_size indépendamment de tp_size. Certains, comme Llama 4 et Gemma 4, utilisent encore par défaut le masquage et l'all-reduce ("ep_router" et "moe_tp_experts" dans model.ep_plan). Le masquage est également disponible sur tout modèle avec une surcharge de ep_plan, et nécessite ep_size=tp_size afin que chaque rang d'un groupe d'experts reçoive les mêmes jetons.

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

Chaque modèle MoE définit deux plans dans sa configuration : base_model_tp_plan pour les modules denses et base_model_ep_plan pour les experts. Ils sont exposés sur le modèle chargé sous la forme de model.tp_plan et model.ep_plan. Avec tp_size > 1 et ep_size > 1, les deux s’appliquent : le plan de parallélisme de tenseurs découpe l’attention et les MLP denses, tandis que le plan de parallélisme d’experts découpe les experts. Les règles EP prennent le pas sur les règles TP pour les mêmes modules, de sorte que les poids des experts ne sont découpés qu’une fois, par le plan EP. Dans le plan EP, le style GroupedGemmParallel répartit les poids des experts le long de la dimension des experts afin que chaque rang ne charge que ses experts locaux, et ep_router masque les experts situés sur d’autres rangs avant qu’un all-reduce ne combine les sorties des experts.

tp_plan n’est appliqué que lorsque tp_size > 1, et ep_plan uniquement lorsque ep_size > 1. Si TP est activé et EP désactivé, le plan TP complet s’applique, y compris les règles d’experts.

La règle de propagation avant d’expert dans ep_plan sélectionne la manière dont les jetons atteignent les experts :

enable_expert_parallel=True sans ep_size est obsolète et émet un FutureWarning. Il ne sharde que les experts, comme ep_size=tp_size, tp_plan={}.

Lancez votre script d'inférence avec torchrun. Le nombre de processus doit être égal à tp_size * fsdp_size * pp_size, et ep_size doit diviser exactement le nombre d'experts.

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

Écraser les plans

Passez tp_plan={...} ou ep_plan={...} à DistributedConfig pour remplacer des règles individuelles des plans prédéfinis. Les règles non spécifiées sont conservées, et les plans fusionnés sont stockés sur le modèle. Chaque clé doit correspondre à un module, à un paramètre ou à une entrée existante du plan ; sinon, le chargement lève une ValueError avant tout fractionnement. Utilisez le chemin complet tel qu’il est vu depuis le modèle chargé, donc model.layers.* pour un LM causal et layers.* pour son modèle de 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"},
)

Fournir un plan n’infère pas les tailles de parallélisme : définissez explicitement tp_size et ep_size.

La plupart des modèles MoE utilisent par défaut "ep_dispatch_experts". Pour utiliser le masquage et l'all-reduce à la place, définissez ep_size=tp_size et écrasez à la fois les règles de propagation avant du routeur et des experts (le nom du module du routeur dépend du modèle, par exemple mlp.router sur 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",
    },
)

À l’inverse, pour un modèle dont le plan utilise le masquage, écrasez la règle de propagation avant des experts avec "ep_dispatch_experts" afin d’utiliser la distribution de jetons. La règle du routeur est alors ignorée, car la distribution a besoin des identifiants globaux des experts pour trouver le propriétaire de chaque expert.

Affectation des jetons

Avec la distribution de jetons, chaque rang entraîne son propre segment du lot. À chaque couche MoE, un routage des jetons est effectué par le rang, chaque paire (jeton, expert) est envoyée au rang propriétaire de l’expert concerné via un all-to-all, les experts locaux sont exécutés sur les données reçues, les résultats sont récupérés par un second all-to-all puis combinés avec les poids de routage. Seules les activations routées et les sorties des experts transitent, et aucun rang ne calcule les experts pour des jetons qui ne lui appartiennent pas.

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)

Avec tp_size=1, ep_size doit diviser fsdp_size et le nombre d'experts, et l'attention n'est pas limitée par num_key_value_heads. Pour le reste du modèle :

• Les paramètres situés en dehors des experts sont fractionnés avec FSDP2 sur fsdp, ce qui réduit leurs gradients.
• Les experts sont répartis sur ep et, lorsque efsdp_size = fsdp_size * tp_size // ep_size est supérieur à un, ils sont également répartis via FSDP sur efsdp. Ils sont toujours enveloppés par FSDP, de sorte que fsdp_mixed_precision et fsdp_cpu_offload s’appliquent également à eux et que save_pretrained() les regroupe comme n’importe quel autre paramètre.
• Un groupe d’expert parallèle détient ep_size / tp_size lots, de sorte que le gradient d’un expert est une somme sur ce nombre de lots. La réduction efsdp divise par fsdp_size au lieu de la taille de son groupe, ce qui donne la même moyenne par lot que celle utilisée par FSDP2 pour les modules denses.
• Chaque expert local traite également une ligne de zéros par couche. Un rang dont les experts n’ont reçu aucun jeton participe toujours à l’opération inverse all-to-all et à la réduction des gradients des experts.

Comment les tailles se combinent

• tp_size * fsdp_size est le nombre de processus. ep_size n’en ajoute aucun : il regroupe les mêmes rangs uniquement pour les poids des experts.
• ep_size découpe la liste des experts en ep_size blocs. Chaque rang calcule num_experts / ep_size experts, et ep_size rangs consécutifs détiennent un ensemble complet. Cet ensemble de rangs constitue le groupe dans lequel s’exécute l’opération all-to-all.
• efsdp_size = fsdp_size * tp_size / ep_size correspond au nombre de copies complètes de l’ensemble d’experts. Les rangs occupant la même position dans les différentes copies découpent ces experts pour la mémoire et moyennent leurs gradients, de la même manière que FSDP le fait pour les modules denses.
• Le lot détenu par un rang ne dépend que de fsdp. Les rangs consécutifs forment un groupe TP et reçoivent le même lot ; les fsdp_size groupes reçoivent des lots différents.

Huit processus, tp_size=2, fsdp_size=4, huit experts :

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

Deux nombres découlent de l'image :

• Au sein d'un groupe EP, chaque jeton existe tp_size fois, une fois par rang de la paire qui détient son lot. Cela ne dépend pas de ep_size.
• Un groupe EP détient ep_size / tp_size lots différents. C'est le nombre sur lequel la gradient d'un expert est additionné.

Avec le parallélisme de tenseurs

Définissez tp_size > 1 pour sharder les modules denses avec le plan TP tandis que les experts utilisent le dispatch. Sur huit processus :

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

Chaque paire de rangs TP reçoit le même lot, car le parallélisme de tenseurs réplique les activations au sein de la paire. Si les deux rangs transmettaient tous leurs jetons, le rang propriétaire d'un expert recevrait chaque jeton deux fois, le calculerait deux fois, et le gradient de ses poids serait doublé. Les experts sont entiers sur un seul rang, de sorte que le doublon ne peut pas être réparti par les poids, et le propriétaire est généralement un autre rang, ce qui empêche de résoudre le problème par la propriété, contrairement au masquage. La paire divise donc les lignes : chaque rang TP transmet une fraction disjointe de 1 / tp_size des jetons, les résultats reviennent au rang qui les a envoyés, et une réduction globale sur la paire des moitiés complétées par des zéros restaure la sortie répliquée attendue par la couche suivante. La division est effectuée par tp_size, et non par ep_size, car seuls les rangs qui détiennent un lot peuvent le transmettre. Les groupes d'experts s'étendent sur quatre rangs et chaque expert est fragmenté par FSDP sur efsdp_size = 2 rangs, tandis que le groupe FSDP du tronc s'étend sur quatre rangs. Les contraintes TP habituelles du modèle, telles que la divisibilité des têtes d'attention, s'appliquent toujours aux modules denses. Les tranches de jetons peuvent être inégales ou vides, y compris lors du décodage mono-jeton.

La distribution de jetons ne peut pas encore être combinée avec le parallélisme de pipeline ; utilisez pp_size=1 (non testé pour l'instant) Ces configurations utilisent chacune huit GPU :

Le Trainer ne prend pas encore en charge la distribution des jetons : le comptage des lots et des jetons suppose la disposition all-reduce, où les rangs d’un groupe TP partagent un lot et où fsdp_size fragments de parallélisme de données existent. La prise en charge de la distribution par le Trainer sera apportée dans une mise à jour ultérieure.

Combinaison avec FSDP2

Le parallélisme de tenseurs et d’experts répartit les poids sur tp, mais l’état de l’optimiseur et les modules sans règle restent répliqués sur chaque rang du groupe, ce qui limite la taille des modèles que vous pouvez entraîner. Ajoutez FSDP2 sur une seconde dimension de maillage avec fsdp_size. Grâce au masquage et à l’all-reduce, conservez ep_size=tp_size et passez `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)

Chargez le modèle comme d’habitude, puis entraînez-le avec Trainer. Il calcule la norme du gradient sur les deux maillages et attribue à chaque maillage son propre groupe de paramètres d’optimiseur. save_model() regroupe les poids fragmentés en un point de contrôle standard. Cela nécessite accelerate>=1.12 afin que Trainer puisse refléter tp_size et fsdp_size dans ~Accelerate.ParallelismConfig.

Le tableau ci-dessous compare l’entraînement EP seul avec l’entraînement 2D EP+FSDP2 sur 8 GPU H100. La charge de travail est le fine-tuning complet de Qwen3-30B-A3B en bf16 avec une longueur de séquence de 2048. Plus il y a de shards FSDP, plus la mémoire de pointe est réduite, et les tokens/s baissent légèrement car FSDP2 effectue des all-gathers et des reduce-scatters des experts à travers fsdp.

La reprise à partir d'un point de contrôle n'est pas encore prise en charge pour les modèles fragmentés au chargement, de sorte que le Trainer n'accepte que save_only_model=True ou save_strategy="no" pour ces modèles.

Référence de l'API