Paralelismo de tensores para el entrenamiento

El paralelismo de tensores (TP) divide las matrices de pesos por columnas o por filas entre varias GPU. Cada GPU almacena un fragmento, calcula un resultado parcial y se sincroniza mediante un all-reduce para producir la salida completa.

TP depende de una comunicación frecuente entre GPUs. Funciona mejor en hardware con enlaces intra-nodo rápidos, como NVLink.

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

Transformers admite TP para arquitecturas cuya configuración define base_model_tp_plan. Comprueba primero ese campo para ver si un modelo admite TP nativo.

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 modelo admite TP, cree una DistributedConfig con el número de dispositivos en tp_size y pásela a from_pretrained(). Transformers utiliza el plan predefinido del modelo, inicializa la malla de dispositivos y fragmenta las capas compatibles por usted.

También puede establecer tp_plan="auto" en DistributedConfig. Cuando se omite tp_size, se infiere a partir de WORLD_SIZE. Pasar tp_plan directamente a from_pretrained() está en desuso y se eliminará en v5.18.

No utilice device_map junto con distributed_config. Ambos entran en conflicto a nivel de carga de pesos. device_map coloca módulos completos en GPUs específicas, mientras que la paralelización de tensores fragmenta esos mismos parámetros entre todas las GPUs.

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 detecta el plan de paralelismo de tensores, lee tp_size del modelo y crea automáticamente una ParallelismConfig.

Inicie el entrenamiento en un nodo con 4 GPUs.

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

ParallelismConfig

Pase ParallelismConfig explícitamente al combinar TP con otras técnicas de paralelismo como 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,
)

Adaptadores PEFT

Puedes cargar adaptadores PEFT en un modelo base con paralelismo de tensores (requiere peft >= 0.21.1). PEFT se encarga de la partición de los pesos del adaptador. Consulta PEFT para saber cómo añadir, cargar y entrenar adaptadores, incluidas las notas sobre entrenamiento distribuido.

Próximos pasos

Lea el capítulo Tensor Parallelism de The Ultra-Scale Playbook para obtener más detalles sobre su funcionamiento.
Lea la guía de inferencia con paralelismo de tensores para obtener más información sobre las estrategias de partición, los planes de TP manuales y los detalles de implementación.

Update on GitHub