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. Les modèles MoE partagent leurs experts avec "ep_dispatch_experts". Ses implémentations déterminent comment les jetons atteignent les experts. Token dispatch permet à ep_size de différer de tp_size. Quelques implémentations EP exigent ep_size=tp_size, par exemple "all_reduce" et "deepgemm_megamoe", de sorte que chaque rang d’un groupe d’experts détient les mêmes jetons.

Llama 4 n’a pas de plan EP : ses experts traitent chaque jeton déjà mis à l’échelle par son score de routeur, il n’y a donc rien à dispatcher. Utilisez plutôt le parallélisme de tenseurs (tp_size) pour les sharder.

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 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 la règle "ep_dispatch_experts" sur le module des experts détermine comment les jetons y parviennent.

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.

"ep_dispatch_experts" répartit des experts entiers par rang (ep_size est un multiple de tp_size qui divise fsdp_size * tp_size) et y accède via l’implémentation nommée par ep_implementation, voir Choisir une implémentation.

Choisir une implémentation

"ep_dispatch_experts" découpe les experts une seule fois et détermine comment les jetons y accèdent à partir des implémentations enregistrées dans EPImplementationInterface, sélectionnées via ep_implementation dans DistributedConfig :

ep_implementation ne choisit que parmi les tailles de parallélisme définies dans DistributedConfig. Lorsque les rangs d’un groupe d’experts détiennent des lots différents, par exemple tp_size=1 avec fsdp_size=ep_size, "all_to_all" est la seule implémentation applicable.

Ajouter une implémentation

Une implémentation sous-classe EPImplementation et s’enregistre sous le nom ep_implementation selects. Elle est construite une fois par module d’experts, ce qui lui permet de conserver un état tel que des tampons de communication. L’appeler avec le lot détenu par ce rang, avec les identifiants globaux des experts, renvoie la sortie pondérée des experts pour chacun de ses jetons. Elle exécute les experts de ce rang avec experts_forward, sur les lignes, les identifiants locaux des experts et les poids qu’elle leur a transmis.

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

Les gradients doivent atteindre chaque rang qui détient un jeton, y compris les poids de routage. Il faut donc envelopper une opération collective que autograd ne différencie pas dans un torch.autograd.Function.

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.

Avec ep_size=tp_size, l'implémentation par défaut "auto" exécute "all_reduce" lors de l'inférence et distribue les jetons lors de l'entraînement. Pour masquer et effectuer un all-reduce également lors de l'entraînement, définissez ep_implementation="all_reduce" :

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

Affectation des jetons

La distribution des jetons est l’implémentation "all_to_all" : chaque rang entraîne sur sa propre partie du lot. À chaque couche MoE, un rang achemine ses jetons, envoie chaque paire (jeton, expert) au rang qui possède l’expert concerné via un all-to-all, exécute ses experts locaux sur ce qu’il reçoit, récupère les résultats avec un second all-to-all et les combine avec les poids d’acheminement. Seules les activations acheminées et les sorties des experts transitent, et aucun rang ne calcule les experts pour des jetons qu’il ne possède 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 le parallélisme d’experts fractionnent les poids sur leurs propres maillages, tp et ep, 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 autre dimension de maillage avec fsdp_size. Avec le masquage et l’all-reduce, conservez ep_size=tp_size et passez 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)

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