Параллелизм тензоров для обучения

Тензорное параллелизм (TP) разделяет матрицы весов по столбцам или по строкам между GPU. Каждый GPU хранит свой фрагмент, вычисляет частичный результат и синхронизируется с помощью all-reduce для получения полного вывода.

TP опирается на частую меж-GPU коммуникацию. Он лучше всего работает на оборудовании с быстрыми внутриузловыми соединениями, такими как NVLink.

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

Transformers поддерживает TP для архитектур, в конфигурации которых определено поле base_model_tp_plan. Сначала проверьте это поле, чтобы убедиться, что модель поддерживает нативный TP.

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)

Если модель поддерживает TP, создайте DistributedConfig с количеством устройств в tp_size и передайте его в from_pretrained(). Transformers использует предопределённый план модели, инициализирует сетку устройств и разбивает поддерживаемые слои на части за вас.

Также можно установить tp_plan="auto" в DistributedConfig. Если tp_size не указан, он вычисляется из WORLD_SIZE. Передача tp_plan непосредственно в from_pretrained() устарела и будет удалена в v5.18.

Не используйте device_map вместе с distributed_config. Эти параметры конфликтуют на уровне загрузки весов. device_map размещает целые модули на конкретных GPU, тогда как тензорное параллелизм распределяет те же самые параметры по всем 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 определяет план тензорного параллелизма, считывает tp_size из модели и автоматически создаёт ParallelismConfig.

Запустите обучение на одном узле с 4 GPU.

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

ParallelismConfig

Явно передавайте ParallelismConfig при сочетании TP с другими методами параллелизма, такими как 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,
)

Адаптеры PEFT

Вы можете загрузить адаптеры PEFT на базовую модель с тензорным параллелизмом (требуется peft >= 0.21.1). PEFT самостоятельно распределяет веса адаптеров. О том, как добавлять, загружать и обучать адаптеры, включая примечания по распределённому обучению, см. в PEFT.

Дальнейшие шаги

Подробнее о том, как это работает, читайте в главе параллелизм тензоров из The Ultra-Scale Playbook.
Ознакомьтесь с руководством по инференсу с тензорным параллелизмом, чтобы узнать больше о стратегиях разделения, ручных планах TP и деталях реализации.

Update on GitHub