추론을 위한 텐서 병렬성

Tensor parallelism은 모델 레이어를 여러 조각으로 나누어 여러 하드웨어 가속기가 동시에 처리하도록 합니다. 이를 통해 단일 GPU의 메모리 용량을 초과하는 모델을 실행하고 더 높은 처리량을 달성할 수 있습니다. 각 레이어에서 GPU가 부분 결과를 교환하므로 노드 내 빠른 통신이 필요합니다.

모델의 config에 base_model_tp_plan이 정의되어 있으면 해당 모델은 tensor parallelism을 지원합니다. 로드된 모델은 supports_tp_plan 속성을 통해 확인할 수 있습니다.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

이 가이드에서는 Transformers에서 텐서 병렬화를 활성화하는 방법과 사용 가능한 분할 전략을 다룹니다.

모델 파티셔닝

DistributedConfig에서 tp_size로 텐서 병렬 장치 수를 설정합니다.

• 모델의 사전 정의된 플랜을 사용하려면 DistributedConfig(tp_size=N)를 설정하세요.
• 수동으로 tp_plan을 정의하고 tp_size와 함께 DistributedConfig에 전달합니다.

tp_plan="auto"를 설정하여 사전 정의된 플랜을 명시적으로 요청할 수도 있습니다. tp_size가 생략되고 tp_plan이 설정된 경우, tp_size는 WORLD_SIZE를 다른 병렬 크기들로 나눈 값에서 유도됩니다. tp_plan을 from_pretrained()에 직접 전달하는 방식은 비추천(deprecated)되며 v5.18에서 제거될 예정입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)
print(model.tp_plan)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

# distributed run
outputs = model(inputs)

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

from transformers import AutoModelForCausalLM, DistributedConfig

tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
    ...
}

distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype="auto",
    distributed_config=distributed_config,
)
print(model.tp_plan)

분할 전략

ALL_PARALLEL_STYLES는 distributed/parallel_utils.py에 정의된 ParallelInterface로, 플랜에서 사용할 수 있는 각 전략 이름을 구성된 전략 인스턴스에 매핑합니다. 모델을 샤딩하는 데 직접적으로 상호작용하지는 않지만, 사용 가능한 이름의 공식 목록입니다. distributed/tensor_parallel.py는 텐서 병렬 전략을 여기에 등록합니다:

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

해당 파일의 모든 tensor parallel 전략은 TensorParallelLayer의 하위 클래스입니다. ep_dispatch_experts와 같은 expert parallel 전략은 distributed/expert_parallel.py에 위치하며, 동일한 훅을 가진 ExpertParallelLayer를 하위 클래스로 상속하고 동일한 ALL_PARALLEL_STYLES에 등록됩니다. 가장 자주 사용하는 ColwiseParallel와 RowwiseParallel은 input_layouts와 output_layouts 배치(placement)를 취하며, 이를 통해 하나의 클래스가 여러 plan 이름을 커버합니다. colwise는 출력을 마지막 차원(dimension) 기준으로 샤드(shard) 상태로 유지하여 후속 rowwise 레이어가 소비하도록 하고, colwise_gather_output는 all-gather를 통해 전체 텐서로 다시 모읍니다.

아래 표에서 각 전략을 설명합니다.

패킹 전략

Weight packing은 여러 linear layer를 하나의 더 큰 layer로 결합합니다. PackedColwiseParallel 및 PackedRowwiseParallel 전략은 packed weights를 올바르게 sharding합니다. 기본 ColwiseParallel 또는 RowwiseParallel 전략은 packed weights를 올바르게 sharding하지 않습니다.

아래 예시에서는 up_proj와 gate_proj를 단일 gate_up_proj 모듈로 묶고, gate_up_proj를 샤딩하려면 packed_rowwise 전략이 필요합니다.

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

forward 패스에서 배치 행렬 곱셈을 사용하여 gate_up_proj 모듈의 출력을 계산합니다.

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

일반적인 Shard는 해당 차원을 연속된 블록으로 나누므로, rank는 각 차원의 슬라이스 대신 gate의 끝부분과 up의 앞부분을 받게 됩니다. packed 전략은 _StridedShard를 split_factor(기본값 2)와 함께 사용하여 분할을 교차(interleave) 처리함으로써, 모든 rank가 두 절반에 해당하는 일치하는 슬라이스를 받도록 하고, 이를 통해 샤딩 후에도 chunk가 여전히 정렬되도록 합니다.

커스텀 파티셔닝 전략

TensorParallelLayer를 distributed/tensor_parallel.py에서 상속하여 커스텀 파티셔닝 전략을 만듭니다. 모든 훅에는 no-op 기본값이 있으므로, 전략에 필요한 훅만 오버라이드하세요.

아래 예제는 ColwiseParallel의 축소된 버전을 살펴봅니다.

1. TensorParallelLayer를 상속하고, 전략이 동작하는 placements를 저장합니다. 베이스 클래스는 __init__를 정의하지 않으므로 super()를 호출할 대상이 없습니다.

   복사됨

   class ColwiseParallel(TensorParallelLayer):
       def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
           self.input_layouts = input_layouts or Replicate()
           self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
           self.use_local_output = use_local_output

2. shard_param을 구현하여 하나의 파라미터를 DTensor 플레이스홀더로 래핑합니다. 메타 텐서에서 실행되므로 distribute_tensor는 메타데이터만 생성하고 데이터를 이동하지 않습니다. 아직 분산할 전체 텐서가 없으므로 src_data_rank=None을 전달합니다.

   복사됨

   def shard_param(self, module, param, mesh):
       meta = module._parameters.get(param)
       if meta is None:
           return
       # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
       placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
       module._parameters[param] = torch.nn.Parameter(
           distribute_tensor(meta, mesh, [placement], src_data_rank=None),
           requires_grad=meta.requires_grad,
       )

3. 입력 및 출력 변환을 구현합니다. install_forward는 모듈의 원래 forward를 감싸서 이를 호출하므로, 텐서를 레이아웃 간에 이동하는 것만 수행하면 됩니다. 열 단위 파티셔닝은 복제된 입력을 기대하고 마지막 차원별로 샤딩된 출력을 생성합니다.

   복사됨

   def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
       x = args[0]
       if not isinstance(x, DTensor):
           x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
       if x.placements != (Replicate(),):
           x = x.redistribute(placements=[Replicate()])
       return (x,) + args[1:], kwargs
   
   def transform_output_post_forward(self, module, output, mesh):
       if not isinstance(output, DTensor):
           output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
       if output.placements != (self.output_layouts,):
           output = output.redistribute(placements=[self.output_layouts])
       return output.to_local() if self.use_local_output else output

   동봉된 ColwiseParallel은 이 위에 빠른 경로를 추가하여, 단순한 nn.Linear 추론 및 로컬 텐서가 필요한 양자화 모듈에 대해 DTensor 왕복을 건너뛰게 합니다. 동일한 처리가 필요한 전략을 사용하는 경우, 복사하기 전에 소스를 읽어 보시기 바랍니다.

4. tp_plan에서 전략을 이름으로 참조할 수 있도록 전략을 등록합니다. 등록 시 전략 인스턴스를 전달하므로, __init__에 전달한 placements가 해당 이름과 함께 저장됩니다. 아래 ()에 주의하세요.

   복사됨

   import torch
   
   from transformers import AutoModelForCausalLM, DistributedConfig
   from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES
   
   ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
   tp_plan = {
       "model.layers.*.self_attn.q_proj": "colwise_custom",
       ...
   }
   distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
   model = AutoModelForCausalLM.from_pretrained(
       "meta-llama/Meta-Llama-3-8B-Instruct",
       dtype=torch.bfloat16,
       distributed_config=distributed_config,
   )

벤치마크

Tensor parallelism은 특히 큰 배치 크기나 긴 시퀀스인 경우 추론 속도를 크게 향상시킵니다.

이 차트는 시퀀스 길이가 512인 Llama에서 단일 forward pass에 대한 예상 속도 향상을 보여줍니다.

설계 구현

Transformers는 프레임워크에 종속되지 않는 방식으로 텐서 병렬성을 구현합니다. torch.distributed의 DeviceMesh와 DTensor에 의존하여 단순하고 확장 가능한 인터페이스를 제공합니다.

DeviceMesh

DeviceMesh는 함께 통신하는 다차원 디바이스 그리드를 생성합니다. 서로 다른 병렬화 전략은 서로 다른 통신 패턴을 요구합니다. 이러한 패턴을 처리하려면 여러 서브 메시를 가진 DeviceMesh를 생성하세요.

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

대부분의 torch.distributed 병렬화 전략은 메시 자체 또는 그 서브 메시에 적용됩니다. 메시는 통신 패턴을 자동으로 처리합니다.

DTensor

DTensor(Distributed Tensor)은 일반적인 텐서 연산 위에 분산 로직을 처리합니다. 텐서 병렬성 하의 모델 가중치는 DTensor로 저장되며, 이를 통해 전략이 명시적인 collective 대신 레이아웃 변경으로 통신을 기술할 수 있습니다.

placements 속성은 DeviceMesh의 장치들에 걸쳐 텐서가 어떻게 배치되는지를 PyTorch에 알려 줍니다. 다음 값을 허용합니다:

• Shard(dimension)은 DeviceMesh에서 생성된 DTensor를 주어진 차원을 기준으로 분할합니다. 열 방향 파티셔닝은 가중치의 출력 특징 차원과 바이어스의 유일한 차원을 샤딩합니다.

  복사됨

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension

  Row-wise partitioning은 bias를 복제하고 입력 특징 차원을 샤딩합니다. bias는 모든 rank에서가 아니라 reduction 이후에 한 번만 추가되기 때문입니다.

  복사됨

  weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
  bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs

• Replicate()는 DTensor를 DeviceMesh에 걸쳐 복제하여 각 장치에 텐서의 전체 사본을 생성합니다.

• Partial()은 텐서가 리덕션을 대기 중임을 나타냅니다. 행 단위 레이어의 forward 출력은 Partial()이며, 이를 Replicate()로 재분배하는 것이 all-reduce를 발생시킵니다.

리소스

• 텐서 병렬성에 대한 자세한 내용은 Ultra-Scale Playbook 섹션에서 확인하세요.

• expert parallelism 가이드를 확인하세요. mixture-of-experts (MoE) 모델을 사용하는 경우, 이러한 모델은 tensor parallelism과 expert parallelism을 지원합니다.

• 텐서 병렬에 대한 빠른 개요를 이해하고, 열 병렬 및 행 병렬 구성이 어떻게 다른지 배우려면 Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand 블로그 게시물을 읽어보세요.

• 학습 환경에서 사용하는 방법을 배우려면 Tensor parallelism 학습 가이드를 참조하세요.