# Parallélisme de tenseurs

[Le parallélisme de tenseurs](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) répartit les matrices de poids d'un modèle sur plusieurs appareils. Chaque appareil détient une tranche par colonne (`"colwise"`) ou par ligne (`"rowwise"`) de chaque couche et calcule un résultat partiel. Un `AllReduce` ou un `AllGather` à la frontière de la couche combine les résultats partiels pour produire la sortie complète. Contrairement au [parallélisme de contexte](./context_parallelism), il réduit la mémoire *de poids* par appareil, ce qui est utile lorsqu'un modèle ne tient pas sur un seul appareil. Reportez-vous à [Inférence distribuée](../training/distributed_inference) pour le comparer aux autres stratégies.

## Charger un modèle fractionné

Passez un [TensorParallelConfig](/docs/diffusers/v0.43.0/fr/api/parallel#diffusers.TensorParallelConfig) à l'argument `parallel_config` de la méthode [from_pretrained()](/docs/diffusers/v0.43.0/fr/api/models/overview#diffusers.ModelMixin.from_pretrained) du modèle. `tp_degree` est le nombre d'appareils sur lesquels effectuer le sharding et doit diviser le nombre de têtes d'attention du modèle. La classe du modèle doit définir un `_tp_plan`, qui associe les noms de modules à un style de sharding. Pour vérifier si un modèle prend en charge le parallélisme de tenseurs, recherchez cet attribut sur sa classe, par exemple [Flux2Transformer2DModel](/docs/diffusers/v0.43.0/fr/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

Pour ajouter un plan à un autre modèle, consultez [Ajouter le parallélisme de tenseurs à un modèle](#add-tensor-parallelism-to-a-model).

Charger de cette manière fragmente le point de contrôle au moment de sa lecture. Chaque rang ne lit que sa propre tranche de chaque poids et le place directement sur son GPU, de sorte qu'aucun rang ne détient jamais un poids de taille complète. La mémoire par rang diminue lorsque `tp_degree` augmente.

```py
import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto its own GPU.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()
```

Enregistrez le script sous le nom `tensor_parallel_flux2.py` et lancez-le avec `torchrun`.

```shell
torchrun --nproc-per-node 4 tensor_parallel_flux2.py
```

Le script définit `tp_degree` sur `world_size`, de sorte que `--nproc-per-node 4` répartit le transformer sur 4 GPU.

Seul le transformer est fragmenté. Chaque rang conserve une copie complète de l'encodeur de texte, de sorte que chaque GPU doit disposer de suffisamment de mémoire pour l'encodeur de texte ainsi que pour sa part du transformer.

## Performance de chargement

Le chargement avec `parallel_config` est plus rapide et utilise moins de mémoire CPU par rang que le chargement du modèle complet suivi de l’appel à `enable_parallelism()`. Le tableau ci-dessous présente les résultats pour [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (transformer uniquement, 32 Md de paramètres, bf16) sur 4x A10G (23 Go).

| tp_degree | méthode | temps de chargement | CPU maximal/rang |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6,8 Go |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64,1 Go |

## Limites

Le parallélisme de tenseurs ne fonctionne pas avec certaines options de chargement et fonctionnalités, et il est impossible d'enregistrer un modèle fractionné.

- Le chargement avec un `parallel_config` parallèle par tenseurs ne fonctionne pas avec `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True` ou des poids non-safetensors. Chacun de ces paramètres provoque une erreur au lieu de charger silencieusement le point de contrôle complet.
- La combinaison de la parallélisation des tenseurs avec la quantification, le déchargement ou les adaptateurs LoRA provoque également une erreur, quelle que soit la méthode de répartition du modèle.
- Pour sharder un modèle déjà en mémoire, appelez `enable_parallelism()` avec la même configuration. Il charge d’abord le modèle complet puis le sharde, ce qui nécessite à chaque rang la mémoire pour le checkpoint complet. Le chargement avec `parallel_config` évite ce coût.
- Enregistrer un modèle parallélisé par tenseurs provoque une erreur dans [save_pretrained()](/docs/diffusers/v0.43.0/fr/api/models/overview#diffusers.ModelMixin.save_pretrained). Enregistrez le modèle avant de le fragmenter.

## Ajouter le parallélisme de tenseurs à un modèle

Cette section s’adresse aux contributeurs souhaitant ajouter la prise en charge du parallélisme de tenseurs à un modèle Diffusers.

Le parallélisme tensoriel ne fonctionne que sur les modèles qui définissent un `_tp_plan`, un attribut de classe qui associe les noms de modules à un style de partitionnement. En rédiger un consiste principalement à associer chaque projection qui *étend* la dimension cachée à la projection qui la *réduit* à nouveau.

Chaque clé peut contenir au plus un `*`, et la partie située avant le `*` doit correspondre à un [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), de sorte qu'une seule entrée couvre tous les blocs. Une clé sans `*` s'applique au modèle lui-même. Les chemins sont relatifs au modèle.

### Colwise et rowwise

| Style | Shards | Chaque rang | Utiliser pour |
|---|---|---|---|
| `"colwise"` | caractéristiques de sortie (`weight` dim 0) | calcule une portion de la sortie | `to_q`, `to_k`, `to_v`, projection d'entrée du FFN |
| `"rowwise"` | caractéristiques d'entrée (dim 1 de `weight`) | calcule une somme partielle | `to_out.0`, projection de sortie du FFN |
| `"rowwise_split_input"` | caractéristiques d'entrée (dim 1 de `weight`) | calcule une somme partielle à partir de sa tranche de l'entrée répliquée | une projection autonome dont l'entrée est répliquée, comme une projection de modulation AdaLN qui lit l'embedding du timestep |

Associez toujours `"colwise"` et `"rowwise"` dans cet ordre. Une projection `"colwise"` laisse sa sortie fragmentée, la projection `"rowwise"` suivante consomme directement ce fragment, et un unique `AllReduce` à la frontière du bloc reconstruit le résultat. Fragmenter la paire de toute autre manière impose un regroupement intermédiaire et communique beaucoup plus.

Dans l'attention, les projections Q, K et V par colonne attribuent à chaque rang un sous-ensemble de têtes, de sorte que `tp_degree` doit diviser le nombre de têtes. Les duplications du flux encodeur (`add_q_proj`, `to_add_out`, `ff_context`) suivent le même schéma que leurs équivalents du flux image.

### Projections fusionnées

Lorsqu’un `Linear` regroupe plusieurs tenseurs logiques le long de la dimension qui est partagée, des tranches simples `"colwise"`/`"rowwise"` traversent la concaténation et désalignent les morceaux. Utilisez plutôt `PackedColwiseParallel`/`PackedRowwiseParallel`, qui partagent chaque bloc regroupé indépendamment.

```py
# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
```

Passer `blocks`, une liste d'entiers proportionnels dont la somme divise la dimension packée. Par exemple, utilisez `[1, 1]` pour une projection de porte et une projection de montée SwiGLU de moitiés égales, ou `[1, 1, 1, 3, 3]` pour une projection fusionnée de Q, K, V, porte et montée avec `mlp_ratio=3`.

```py
"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),
```

Lorsque les tailles de blocs ne sont connues que depuis la configuration, omettez `blocks` et stockez les tailles absolues sur le `Linear` lors de `__init__` à la place, sous la forme de `_tp_packed_col_blocks` ou `_tp_packed_row_blocks`.

```py
# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]
```

```py
# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
```

Le plan dans [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) utilise les deux formes. La paire SwiGLU double-flux transmet directement `blocks`. La projection QKV et MLP fusionnée en flux unique stocke ses tailles sur le `Linear` dans `Flux2ParallelSelfAttention.__init__`, car elles dépendent de `mlp_ratio` et `mlp_mult_factor`.

### À exclure

Les modules exclus du plan sont répliqués sur chaque rang. Excluez les couches de normalisation, les embeddings de patch et de texte, ainsi que les `norm_out` et `proj_out` finaux. Ils sont petits, donc les sharder économise peu de mémoire et ajoute de la communication. La modulation AdaLN (`img_mod` et `txt_mod`) est généralement assez petite pour être exclue également. Lorsqu'elle est grande, comme `adaln_proj.linear` de MiniMax-H3, shardez-la avec `"rowwise_split_input"`.

### Contraintes et vérification

- `tp_degree` doit diviser `config.num_attention_heads`.
- Chaque bloc compacté doit être *individuellement* divisible par `tp_degree`, et non seulement leur somme.

[from_pretrained()](/docs/diffusers/v0.43.0/fr/api/models/overview#diffusers.ModelMixin.from_pretrained) et `enable_parallelism()` effectuent les deux vérifications avant de charger ou de fractionner un poids.

Valider un nouveau plan numériquement plutôt qu’à l’œil : générer avec une graine fixe sur un seul appareil, puis à nouveau sous parallélisme de tenseurs, et comparer les sorties. Un `"colwise"`/`"rowwise"` mal placé s’exécute généralement tout de même et produit une image plausible mais incorrecte.

> [!TIP]
> Commencez à partir d'un plan existant pour une architecture similaire.
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) ne contient pas de projections fusionnées, chaque entrée est donc simplement `"colwise"` ou `"rowwise"`. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) ajoute une projection row-wise packée. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) utilise les deux styles packés.