Parallelisme tensor untuk pelatihan

Tensor parallelism (TP) membagi matriks bobot secara kolom atau baris di antara beberapa GPU. Setiap GPU menyimpan satu shard, menghitung hasil parsial, dan melakukan sinkronisasi melalui all-reduce untuk menghasilkan output lengkap.

TP bergantung pada komunikasi antar-GPU yang sering. TP bekerja paling baik pada perangkat keras dengan tautan intra-node yang cepat, seperti NVLink.

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

Transformers mendukung TP untuk arsitektur yang konfigurasinya mendefinisikan base_model_tp_plan. Periksa field tersebut terlebih dahulu untuk melihat apakah sebuah model mendukung TP native.

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)

Jika sebuah model mendukung TP, buat DistributedConfig dengan jumlah perangkat di tp_size dan teruskan ke from_pretrained(). Transformers menggunakan rencana yang telah didefinisikan sebelumnya oleh model, menginisialisasi device mesh, dan memecah lapisan yang didukung untuk Anda.

Anda juga dapat mengatur tp_plan="auto" di DistributedConfig. Ketika tp_size diabaikan, nilainya disimpulkan dari WORLD_SIZE. Menyalurkan tp_plan secara langsung ke from_pretrained() sudah tidak didukung dan akan dihapus pada v5.18.

Jangan gunakan device_map bersama distributed_config. Keduanya saling bertentangan pada level pemuatan bobot. device_map menempatkan modul utuh pada GPU tertentu, sedangkan tensor parallelism memecah parameter yang sama tersebut di seluruh 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 mendeteksi rencana parallel tensor, membaca tp_size dari model, dan membuat ParallelismConfig secara otomatis.

Luncurkan pelatihan pada satu node dengan 4 GPU.

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

ParallelismConfig

Lewatkan ParallelismConfig secara eksplisit saat menggabungkan TP dengan teknik paralelisasi lain seperti 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,
)

Adaptasi PEFT

Anda dapat memuat adapter PEFT ke model dasar yang menggunakan tensor parallel (memerlukan peft >= 0.21.1). PEFT menangani sharding bobot adapter. Lihat PEFT untuk cara menambahkan, memuat, dan melatih adapter, termasuk catatan mengenai pelatihan terdistribusi.

Langkah selanjutnya

• Baca bab Tensor Parallelism dari The Ultra-Scale Playbook untuk detail lebih lanjut tentang cara kerjanya.
• Baca panduan inferensi parallelisme tensor untuk mempelajari lebih lanjut tentang strategi partisi, rencana TP manual, dan detail implementasi.