Parallélisme de tenseurs pour l'inférence

Le parallélisme de tenseurs découpe une couche de modèle en plusieurs morceaux afin que plusieurs accélérateurs matériels puissent y travailler simultanément. Cela permet d'exécuter des modèles dont la taille dépasse la capacité mémoire d'un seul GPU et d'atteindre un débit supérieur. Une communication intra-nœud rapide est nécessaire, car les GPU échangent des résultats partiels à chaque couche.

Un modèle prend en charge le parallélisme de tenseurs si sa configuration définit base_model_tp_plan. Vérifiez un modèle chargé à l'aide de la propriété supports_tp_plan.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

Ce guide explique comment activer le parallélisme de tenseurs dans Transformers et les stratégies de partitionnement disponibles.

Partitionnement d'un modèle

Configurez le nombre d’appareils de parallélisme de tenseurs avec tp_size dans DistributedConfig.

Définissez DistributedConfig(tp_size=N) pour utiliser le plan prédéfini du modèle.
Définissez un tp_plan manuel et transmettez-le à DistributedConfig avec tp_size.

Vous pouvez également définir tp_plan="auto" pour demander explicitement le plan prédéfini. Lorsque tp_size est omis et qu'un tp_plan est défini, tp_size est dérivé de WORLD_SIZE divisé par les autres tailles de parallélisme. Transmettre tp_plan directement à from_pretrained() est déprécié et sera supprimé dans la v5.18.

plan auto

plan manuel

Stratégies de partitionnement

ALL_PARALLEL_STYLES, une ParallelInterface définie dans distributed/parallel_utils.py, associe chaque nom de stratégie utilisable dans un plan à une instance de stratégie configurée. Vous n’interagissez pas directement avec elle pour sharder un modèle, mais elle constitue la liste officielle des noms disponibles. distributed/tensor_parallel.py y enregistre les stratégies de parallélisme tensoriel :

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

Chaque stratégie de parallélisme tensoriel est une sous-classe de TensorParallelLayer dans ce fichier. Les stratégies de parallélisme d'experts, telles que ep_dispatch_experts, se trouvent dans distributed/expert_parallel.py, héritent de ExpertParallelLayer, qui possède les mêmes hooks, et s'enregistrent dans le même ALL_PARALLEL_STYLES. Les deux que vous utiliserez le plus souvent, ColwiseParallel et RowwiseParallel, prennent des placements input_layouts et output_layouts, ce qui permet à une seule classe de couvrir plusieurs noms de plans. colwise laisse sa sortie fragmentée sur la dernière dimension pour qu'une couche rowwise suivante puisse la consommer, tandis que colwise_gather_output la réassemble en un tenseur complet via un all-gather.

Le tableau ci-dessous décrit chaque stratégie.

Stratégie	Description
`ColwiseParallel`	Fragmente un poids linéaire sur sa dimension de caractéristiques de sortie (`Shard(0)` pour un poids `nn.Linear` 2D) et fragmente un biais 1D. Redistribue l'entrée vers `Replicate()`, puis place la sortie dans `output_layouts`.
`RowwiseParallel`	Fractionne le poids linéaire sur sa dimension de fonction d’entrée (`Shard(-1)`) et réplique le biais, qui est ajouté après la réduction. Le forward de chaque rang produit une sortie `Partial()` contenant sa part de la somme, et `transform_output_post_forward` la redistribue vers `output_layouts`. Réduire vers `Replicate()` déclenche un all-reduce, et réduire vers `Shard(1)` déclenche un reduce-scatter.
`EmbeddingParallel`	Fragmente un `nn.Embedding` en répliquant les identifiants de jetons. `shard_dim=0` (`embedding_rowwise`, le choix habituel) fragmente le vocabulaire et effectue une réduction globale (all-reduce) des résultats partiels de chaque rang pour obtenir la sortie complète. `shard_dim=1` (`embedding_colwise`) fragmente la dimension cachée et effectue une collecte globale (all-gather) des tranches de sortie de chaque rang.
`PackedColwiseParallel`	Une variante de `ColwiseParallel` pour les poids fusionnés, par exemple `up_proj` et `gate_proj` regroupés dans `gate_up_proj`. Utilisez `split_factor` lorsqu'un poids regroupe un nombre de couches différent de deux.
`PackedRowwiseParallel`	L'équivalent par ligne, pour les poids empilés le long de la dernière dimension. Réplique les paramètres 1D.
`SequenceParallel`	Duplique les paramètres du module et sharde son entrée sur `sequence_dim` (par défaut `1`). Utilisé pour les normalisations qui opèrent par jeton, telles que `LayerNorm` et `RMSNorm`.
`ReplicatedWithGradAllReduce`	Duplique un paramètre mais effectue un all-reduce sur son gradient. Nécessaire pour les normes situées entre une couche en colonnes et une couche en lignes, qui normalisent le long d'un axe fragmenté, où chaque rang ne voit que ses propres têtes.
`AllReduceParallel`	Effectue une réduction globale (all-reduce) de la sortie avant `Partial()` d'un module vers `Replicate()`. Utilisez-le comme point de synchronisation pour un module dont le calcul se termine par une somme partielle.
`MlaKvAProjParallel`	Découpe la sortie de `kv_a_proj_with_mqa` de l'attention MLA de type DeepSeek-V2 et effectue un all-reduce sur le gradient de la moitié RoPE, ce qui contourne `kv_b_proj` et éviterait sinon de conserver un gradient partiel. Nécessite `qk_rope_head_dim` dans la configuration du modèle.
`MoeExpertsParallel`	Experts MoE en parallélisme de tenseurs. Effectue un all-reduce sur la sortie avant des experts et ajoute les all-reduces de rétropropagation pour les états cachés et les poids de routage.
`MoeIdentityParallel`	Divise par avance l’entrée d’un expert nul ou identité par la taille du maillage, annulant ainsi l’opération all-reduce appliquée en aval par `moe_tp_experts`.
`LocalParamsParallel`	Exécute le forward d’un module avec ses propres paramètres shardés en tant que tenseurs locaux. À utiliser sur le propriétaire d’un `nn.Parameter` brut, tel que les attention sinks, que sa propre entrée de plan sharde.
`MoEParamShard`	Répartit les poids des experts MoE sur un placement donné. Soutient le nom `grouped_gemm`, où `shards_expert_dim=True` réécrit également `module.num_experts` en le nombre d'experts par rang.
`EpRouterParallel`	Masque les scores du routeur pour les experts non locaux et remappe les identifiants globaux d'experts vers les identifiants locaux, afin que chaque rang n'exécute que les experts qu'il possède. Exige que `num_experts` soit divisible par la taille du maillage.
`RouterParallelMegaMoe`	Variante de routeur pour DeepGEMM Mega MoE, qui achemine les experts à l’intérieur du noyau, ce dernier souhaitant que la sortie du routeur reste inchangée.
`MoeTensorParalellMegaMoeExperts`	Couche d’experts dédiée à l’inférence pour DeepGEMM Mega MoE. Elle omet les synchronisations de gradients et transmet le groupe de processus au module afin que le noyau puisse initialiser ses tampons partagés lors du premier passage avant.

Stratégies compactes

Le regroupement de poids combine plusieurs couches linéaires en une seule couche plus grande. Les stratégies PackedColwiseParallel et PackedRowwiseParallel fractionnent correctement les poids regroupés. Les stratégies de base ColwiseParallel ou RowwiseParallel fractionnent incorrectement les poids regroupés.

L’exemple ci-dessous regroupe up_proj et gate_proj en un seul module gate_up_proj et requiert la stratégie packed_rowwise pour sharder gate_up_proj.

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

Utilisez la multiplication matricielle par lots dans le passage forward pour calculer la sortie du module gate_up_proj.

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

Un Shard simple divise cette dimension en blocs contigus, de sorte qu'un rang reçoit la fin de gate et le début de up au lieu d'une tranche de chacun. Les stratégies empilées utilisent _StridedShard avec split_factor (2 par défaut) pour entrelacer la division, ce qui donne à chaque rang une tranche correspondante des deux moitiés afin que chunk reste aligné après le sharding.

Stratégies de partitionnement personnalisées

Hériter de TensorParallelLayer dans distributed/tensor_parallel.py pour créer une stratégie de partitionnement personnalisée. N’implémenter que les hooks dont votre stratégie a besoin, car chacun d’entre eux possède un comportement par défaut sans effet.

Hook	Objectif
`validate_param`	Rejeter un paramètre que cette stratégie ne peut pas fractionner, avant le chargement de tout poids.
`shard_param`	Remplacez un paramètre par un espace réservé `DTensor` afin que le chargeur sache quelle partition appartient à ce rang.
`transform_inputs_pre_forward`	Redistribuer les entrées du module vers la disposition attendue par son forward.
`context_around_forward`	Enveloppez le forward dans un gestionnaire de contexte, par exemple pour exposer des tenseurs locaux à un noyau.
`transform_output_post_forward`	Redistribuer ou réduire la sortie du module.
`should_use_local_tensors`	Signalez que la méthode forward de ce module nécessite des tenseurs simples plutôt que des `DTensor`.
`install_forward`	Remplacez `module.forward` en totalité. N’implémentez cette méthode que si les hooks ci-dessus ne suffisent pas, comme le fait `ReplicatedWithGradAllReduce` pour enregistrer un hook de rétropropagation.

L'exemple ci-dessous détaille une version simplifiée de ColwiseParallel.

Hériter de TensorParallelLayer et stocker les placements avec lesquels la stratégie fonctionne. La classe de base ne définit pas de __init__, il n’y a donc rien sur lequel appeler super().

class ColwiseParallel(TensorParallelLayer):
    def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
        self.input_layouts = input_layouts or Replicate()
        self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
        self.use_local_output = use_local_output

Implémentez shard_param pour encapsuler un paramètre sous forme de placeholder DTensor. Il s’exécute sur des tenseurs meta, de sorte que distribute_tensor ne construit que les métadonnées et ne déplace aucune donnée. Passez src_data_rank=None car il n’existe pas encore de tenseur complet à disperser.

def shard_param(self, module, param, mesh):
    meta = module._parameters.get(param)
    if meta is None:
        return
    # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
    placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
    module._parameters[param] = torch.nn.Parameter(
        distribute_tensor(meta, mesh, [placement], src_data_rank=None),
        requires_grad=meta.requires_grad,
    )

Implémentez les transformations d’entrée et de sortie. install_forward les appelle autour du forward d’origine du module, de sorte qu’elles n’ont qu’à déplacer les tenseurs entre les différents layouts. Le partitionnement par colonnes s’attend à une entrée répliquée et produit une sortie fragmentée sur la dernière dimension.

def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
    x = args[0]
    if not isinstance(x, DTensor):
        x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
    if x.placements != (Replicate(),):
        x = x.redistribute(placements=[Replicate()])
    return (x,) + args[1:], kwargs

def transform_output_post_forward(self, module, output, mesh):
    if not isinstance(output, DTensor):
        output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
    if output.placements != (self.output_layouts,):
        output = output.redistribute(placements=[self.output_layouts])
    return output.to_local() if self.use_local_output else output

Le ColwiseParallel fourni ajoute des chemins rapides par-dessus ce mécanisme, qui évitent l’aller-retour DTensor pour l’inférence sur des nn.Linear simples et pour les modules quantisés nécessitant des tenseurs locaux. Lisez le code source avant de le copier si votre stratégie nécessite le même traitement.

Enregistrez la stratégie afin qu'un tp_plan puisse la nommer. L'enregistrement prend une instance de stratégie, de sorte que les placements que vous passez à __init__ sont stockés avec le nom. Notez le () ci-dessous.

import torch

from transformers import AutoModelForCausalLM, DistributedConfig
from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES

ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise_custom",
    ...
}
distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

Benchmarks

Le parallélisme de tenseurs accélère considérablement l’inférence, en particulier pour les grandes tailles de lot ou les séquences longues.

Ce graphique montre l’accélération attendue pour un seul passage avant sur Llama avec une longueur de séquence de 512.

Implémentation de la conception

Transformers implémente le parallélisme de tenseurs de manière indépendante du framework. Il s’appuie sur DeviceMesh et DTensor de torch.distributed pour fournir une interface simple et extensible.

DeviceMesh

DeviceMesh crée une grille multi-dimensionnelle d’appareils qui communiquent entre eux. Les différentes stratégies de parallélisation nécessitent des motifs de communication distincts. Créez un DeviceMesh avec plusieurs sous-grilles pour gérer ces motifs.

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

La plupart des stratégies de parallélisation de torch.distributed s’appliquent au maillage lui-même ou à ses sous-maillages. Le maillage gère automatiquement les motifs de communication.

DTensor

DTensor (Distributed Tensor) gère la logique distribuée par-dessus les opérations tensorielles habituelles. Les poids du modèle sous parallélisme de tenseurs sont stockés sous forme de DTensor, ce qui permet à une stratégie de décrire la communication comme un changement de disposition plutôt que comme une opération collective explicite.

L'attribut placements indique à PyTorch comment un tenseur est réparti sur les appareils d'un DeviceMesh. Il accepte les valeurs suivantes :

Shard(dimension) répartit un DTensor sur une dimension donnée, sur le DeviceMesh sous lequel il a été construit. Le partitionnement par colonnes sharde la dimension des caractéristiques de sortie du poids et la seule dimension du biais.
```
weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension
```
Le partitionnement par lignes sharde la dimension des caractéristiques d’entrée et réplique le biais, car le biais est ajouté une seule fois après la réduction, et non sur chaque rang.
```
weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs
```
Replicate() réplique un DTensor sur le DeviceMesh, créant une copie complète du tenseur sur chaque appareil.
Partial() marque un tenseur comme en attente d’une réduction. La sortie avant d’une couche par ligne est Partial(), et sa redistribution vers Replicate() déclenche l’all-reduce.

Ressources

La section Ultra-Scale Playbook sur le parallélisme de tenseurs fournit plus de détails.
Consultez le guide sur le expert parallelism si vous utilisez un modèle de type mixture-of-experts (MoE). Ces modèles prennent en charge le tensor parallelism et l'expert parallelism.
Lisez l’article de blog Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand pour un aperçu rapide du parallélisme de tenseurs et découvrez comment les configurations parallèles par colonne et par ligne diffèrent.
Consultez le guide d’entraînement sur le Tensor parallelism pour apprendre à l’utiliser dans un contexte d’entraînement.

Update on GitHub