Song song tensor cho huấn luyện

Song song tensor (TP) chia các ma trận trọng số theo cột hoặc theo hàng trên các GPU. Mỗi GPU giữ một phần, tính toán một kết quả một phần, và đồng bộ hóa bằng all-reduce để tạo ra đầu ra đầy đủ.

TP dựa vào việc liên lạc liên GPU thường xuyên. Nó hoạt động tốt nhất trên phần cứng có liên kết nội node nhanh như NVLink.

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

Transformers hỗ trợ TP cho các kiến trúc có config định nghĩa base_model_tp_plan. Hãy kiểm tra trường này trước để xem một model có hỗ trợ TP native hay không.

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)

Nếu một model hỗ trợ TP, hãy tạo một DistributedConfig với số lượng thiết bị trong tp_size và truyền nó vào from_pretrained(). Transformers sẽ sử dụng kế hoạch được định nghĩa trước của model, khởi tạo device mesh và chia nhỏ các layer được hỗ trợ cho bạn.

Bạn cũng có thể đặt tp_plan="auto" trong DistributedConfig. Khi tp_size bị bỏ qua, nó sẽ được suy ra từ WORLD_SIZE. Việc truyền tp_plan trực tiếp vào from_pretrained() đã bị ngừng hỗ trợ và sẽ bị loại bỏ trong v5.18.

Không sử dụng device_map cùng với distributed_config. Hai cấu hình này xung đột ở cấp độ nạp trọng số. device_map đặt các module hoàn chỉnh lên các GPU cụ thể, trong khi tensor parallelism chia nhỏ chính các tham số đó trên tất cả các 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 phát hiện kế hoạch song song tensor, đọc tp_size từ mô hình và tự động tạo ParallelismConfig.

Khởi động huấn luyện trên một nút với 4 GPU.

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

ParallelismConfig

Truyền ParallelismConfig một cách tường minh khi kết hợp TP với các kỹ thuật song song khác như 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,
)

Bộ thích ứng PEFT

Bạn có thể nạp các adapter PEFT lên một mô hình cơ sở song song tensor (yêu cầu peft >= 0.21.1). PEFT xử lý việc chia nhỏ trọng số adapter. Xem PEFT để biết cách thêm, nạp và huấn luyện adapter, bao gồm cả các ghi chú về huấn luyện phân tán.

Các bước tiếp theo

Đọc chương Tensor Parallelism trong The Ultra-Scale Playbook để biết thêm chi tiết về cách nó hoạt động.
Đọc hướng dẫn suy luận song song tensor để tìm hiểu thêm về các chiến lược phân vùng, kế hoạch TP thủ công và các chi tiết triển khai.

Update on GitHub