Paralelismo de tensor para treinamento

O paralelismo de tensores (TP) divide as matrizes de pesos por colunas ou por linhas entre as GPUs. Cada GPU mantém um fragmento, calcula um resultado parcial e sincroniza com um all-reduce para produzir a saída completa.

O TP depende de comunicação frequente entre GPUs. Ele funciona melhor em hardware com links intra-nó rápidos, como o NVLink.

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

O Transformers oferece suporte a TP para arquiteturas cujo config define base_model_tp_plan. Verifique esse campo primeiro para ver se um modelo oferece suporte a 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)

Se um modelo suporta TP, crie um DistributedConfig com o número de dispositivos em tp_size e passe-o para from_pretrained(). O Transformers usa o plano predefinido do modelo, inicializa a malha de dispositivos e faz o sharding das camadas suportadas para você.

Você também pode definir tp_plan="auto" em DistributedConfig. Quando tp_size é omitido, ele é inferido a partir de WORLD_SIZE. Passar tp_plan diretamente para from_pretrained() está obsoleto e será removido na v5.18.

Não use device_map com distributed_config. Os dois conflitam no nível de carregamento de pesos. device_map posiciona módulos inteiros em GPUs específicas, enquanto a paralelização de tensores divide esses mesmos parâmetros entre todas as 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 o plano de paralelismo de tensor, lê tp_size do modelo e cria automaticamente um ParallelismConfig.

Inicie o treinamento em um nó com 4 GPUs.

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

ParallelismConfig

Passe ParallelismConfig explicitamente ao combinar TP com outras 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

Você pode carregar adaptadores PEFT em um modelo base com paralelismo de tensor (requer peft >= 0.21.1). O PEFT gerencia o particionamento dos pesos do adaptador. Consulte PEFT para saber como adicionar, carregar e treinar adaptadores, incluindo as observações sobre treinamento distribuído.

Próximas etapas

Leia o capítulo Tensor Parallelism do The Ultra-Scale Playbook para mais detalhes sobre como ele funciona.
Leia o guia de inferência com paralelismo de tensor para saber mais sobre estratégias de particionamento, planos de TP manuais e detalhes de implementação.

Update on GitHub