Parallélisme de tenseurs

Le parallélisme de tenseurs 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, 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 pour le comparer aux autres stratégies.

Charger un modèle fractionné

Passez un TensorParallelConfig à l’argument parallel_config de la méthode from_pretrained() du modèle. tp_degree est le nombre d’appareils sur lesquels effectuer le fractionnement 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 fractionnement. Pour vérifier si un modèle prend en charge le parallélisme de tenseurs, recherchez cet attribut sur sa classe, par exemple Flux2Transformer2DModel.

Pour ajouter un plan à un autre modèle, consultez Ajouter le parallélisme de tenseurs à un modèle.

Chargé de cette manière, le point de contrôle est fragmenté à la lecture. Chaque rang ne lit que sa propre portion 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.

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.

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 (transformer uniquement, 32 Md de paramètres, bf16) sur 4x A10G (23 Go).

Limites

Le parallélisme de tenseurs ne fonctionne pas avec certaines options et fonctionnalités de chargement, 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(). 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, 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

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.

# 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.

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

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

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

Le plan dans 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() 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.

Commencez à partir d'un plan existant pour une architecture similaire.

• transformer_qwenimage.py ne contient pas de projections fusionnées, chaque entrée est donc simplement "colwise" ou "rowwise".
• transformer_flux.py ajoute une projection par ligne empilée.
• transformer_flux2.py utilise les deux styles empilés.