Parallélisme de tenseurs pour l'entraînement

Le parallélisme de tenseurs (TP) répartit les matrices de poids par colonnes ou par lignes sur plusieurs GPU. Chaque GPU conserve un fragment, calcule un résultat partiel et se synchronise via une opération all-reduce pour produire la sortie complète.

TP repose sur des communications fréquentes entre GPU. Il fonctionne de manière optimale sur du matériel doté de liaisons intra-nœud rapides, telles que NVLink.

    ┌─────────────────────────────┐
    │       X  (replicated)       │
    └────┬──────────┬─────────┬───┘
         │          │         │
    ┌────▼───┐ ┌────▼───┐ ┌───▼────┐
    │ ▓▓▓ W₀ │ │ ░░░ W₁ │ │ ███ W₂ │
    │  X@W₀  │ │  X@W₁  │ │  X@W₂  │
    └────┬───┘ └────┬───┘ └───┬────┘
         └──────────┼─────────┘
               Y₀+Y₁+Y₂
    ┌────────────────────────────┐
    │          Y (full)          │
    └────────────────────────────┘

Transformers prend en charge le TP pour les architectures dont la configuration définit base_model_tp_plan. Vérifiez d'abord ce champ pour savoir si un modèle prend en charge le TP natif.

from transformers import AutoConfig

config = AutoConfig.from_pretrained("Qwen/Qwen3-0.6B")
print(config.base_model_tp_plan is not None)
print(config.base_model_tp_plan)

Si un modèle prend en charge le TP, créez un DistributedConfig avec le nombre d’appareils dans tp_size et transmettez-le à from_pretrained(). Transformers utilise le plan prédéfini du modèle, initialise le maillage d’appareils et fragmente les couches prises en charge pour vous.

Vous pouvez également définir tp_plan="auto" dans DistributedConfig. Lorsque tp_size est omis, il est déduit de WORLD_SIZE. Transmettre tp_plan directement à from_pretrained() est déprécié et sera supprimé dans la v5.18.

N’utilisez pas device_map avec distributed_config. Les deux sont en conflit au niveau du chargement des poids. device_map place des modules entiers sur des GPU spécifiques, tandis que le parallélisme de tenseurs répartit ces mêmes paramètres sur l’ensemble des GPU.

import torch

from transformers import AutoModelForCausalLM, DistributedConfig

distributed_config = DistributedConfig(tp_size=4)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

Trainer détecte le plan de parallélisme de tenseurs, lit tp_size depuis le modèle et crée automatiquement un ParallelismConfig.

Lancer l’entraînement sur un nœud avec 4 GPU.

torchrun --nproc-per-node 4 train_tp.py

ParallelismConfig

Transmettez ParallelismConfig explicitement lors de la combinaison de TP avec d'autres techniques de parallélisme telles que FSDP.

import torch

from accelerate import ParallelismConfig
from transformers import AutoModelForCausalLM, DistributedConfig, TrainingArguments

distributed_config = DistributedConfig(tp_size=4)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

parallelism_config = ParallelismConfig(tp_size=4)

args = TrainingArguments(
    ...,
    parallelism_config=parallelism_config,
)

Adaptateurs PEFT

Vous pouvez charger des adaptateurs PEFT sur un modèle de base parallélisé par tenseurs (nécessite peft >= 0.21.1). PEFT gère le fractionnement des poids des adaptateurs. Consultez PEFT pour savoir comment ajouter, charger et entraîner des adaptateurs, y compris les notes relatives à l’entraînement distribué.

Étapes suivantes

Consultez le chapitre Tensor Parallelism du The Ultra-Scale Playbook pour plus de détails sur son fonctionnement.
Consultez le guide d’inférence avec parallélisme de tenseurs pour en savoir plus sur les stratégies de partitionnement, les plans TP manuels et les détails d’implémentation.

Update on GitHub