학습을 위한 텐서 병렬

Tensor parallelism(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는 base_model_tp_plan을 정의하는 config를 가진 아키텍처에 대해 TP를 지원합니다. 모델이 네이티브 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를 지원하는 경우, tp_size에 장치 수를 지정하여 DistributedConfig를 생성하고 from_pretrained()에 전달합니다. Transformers는 모델의 사전 정의된 플랜을 사용하고, device mesh를 초기화하며, 지원되는 레이어를 자동으로 샤딩합니다.

tp_plan="auto"를 DistributedConfig에 설정할 수도 있습니다. tp_size가 생략된 경우 WORLD_SIZE에서 추론됩니다. tp_plan을 from_pretrained()에 직접 전달하는 방식은 비추천(deprecated)되며 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가 있는 1개의 노드에서 학습을 시작합니다.

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

ParallelismConfig

ParallelismConfig을 FSDP와 같은 다른 병렬화 기법과 TP를 조합할 때 명시적으로 전달하세요.

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,
)

PEPT 어댑터

텐서 병렬 기반 모델에 PEFT 어댑터를 로드할 수 있습니다(peft >= 0.21.1 필요). PEFT가 어댑터 가중치의 샤딩을 처리합니다. 어댑터를 추가, 로드 및 학습하는 방법과 분산 학습 관련 참고 사항은 PEFT를 참조하세요.

다음 단계

동작 방식에 대한 자세한 내용은 The Ultra-Scale Playbook의 텐서 병렬 처리 장을 참고하세요.
텐서 병렬 추론 가이드를 읽어 분할 전략, 수동 TP 계획, 구현 세부 사항에 대해 자세히 알아보세요.

Update on GitHub