추론용 텐서 병렬화

추론용 텐서 병렬화 (Tensor parallelism for inference)

텐서 병렬화는 모델 레이어를 조각으로 잘라 여러 하드웨어 가속기가 그것을 동시에 처리하게 합니다. 이를 통해 단일 GPU의 메모리 용량을 넘는 모델을 실행하고 더 높은 처리량을 얻을 수 있어요. GPU가 각 레이어에서 부분 결과를 교환하기 때문에 빠른 노드 내(intra-node) 통신이 필요합니다.

출처: 문서

본문

모델의 설정이 base_model_tp_plan을 정의하면 그 모델은 텐서 병렬화를 지원합니다. 로드된 모델에서 supports_tp_plan 프로퍼티로 확인해 보세요.

from transformers import AutoModelForCausalLM

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

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

모델 분할하기 (Partitioning a model)

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으로 추론 스크립트를 실행합니다. GPU당 프로세스 하나를 사용하세요.

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

각 레이어에 대한 텐서 병렬 플랜을 tp_plan에 정의하고 DistributedConfig를 통해 넘깁니다. 아래 예시는 열(column)과 행(row) 분할을 사용합니다. 지원되는 다른 전략은 Partitioning strategies 섹션을 참고하세요.

수동 분할은 모델 아키텍처와 전략 상호작용에 대한 깊은 이해가 필요해요. 잘못된 분할 선택은 실패하거나 잘못된 결과를 만드는 느린 모델을 만들어냅니다. Ultra-Scale Playbook이 분할 전략을 자세히 설명해요.

키는 모듈 또는 파라미터 이름이고, *는 레이어 인덱스를 대신합니다. 인식되지 않는 전략 이름은 지원되는 이름을 나열한 ValueError를 발생시키며, Transformers는 어떤 규칙과도 일치하지 않은 플랜 규칙과 어떤 규칙으로도 덮이지 않은 파라미터에 대해 경고를 로그합니다.

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)

분할 전략 (Partitioning strategies)

ParallelInterface 클래스는 tp_plan에서 사용할 수 있는 각 전략 이름을 설정된 전략 인스턴스에 매핑합니다. 모델을 분할하기 위해 직접 상호작용하지는 않지만, 사용 가능한 이름의 권위 있는 목록이에요.

class ParallelInterface(GeneralInterface):
    _global_mapping = {
        "embedding_rowwise": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
        "colwise_gather_output": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
        "colwise_rep": ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
        "colwise": ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)),
        "rowwise": RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()),
        "rowwise_split_input": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
        "rowwise_rep": RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()),
        "packed_colwise": PackedColwiseParallel(),
        "packed_rowwise": PackedRowwiseParallel(),
        "sequence_parallel": SequenceParallel(use_local_output=True),
        "grouped_gemm": MoEParamShard(Shard(0), shards_expert_dim=True),
        "ep_router": EpRouterParallel(),
        "megamoe_router": RouterParallelMegaMoe(),
        "moe_tp_experts": MoeExpertsParallel(),
        "megamoe_experts": MoeTensorParalellMegaMoeExperts(),
        "moe_identity_expert": MoeIdentityParallel(),
        "replicated_with_grad_allreduce": ReplicatedWithGradAllReduce(),
        "mla_kv_a_proj": MlaKvAProjParallel(),
        "all_reduce": AllReduceParallel(),
    }

모든 전략은 distributed/tensor_parallel.py의 TensorParallelLayer 하위 클래스입니다. 가장 많이 쓰게 될 두 가지인 ColwiseParallel과 RowwiseParallel은 input_layouts와 output_layouts 배치를 받는데, 그래서 하나의 클래스가 여러 플랜 이름을 다룰 수 있어요. colwise는 출력을 마지막 차원에서 샤딩된 채로 남겨 다음 rowwise 레이어가 소비하게 하며, colwise_gather_output은 it을 다시 전체 텐서로 all-gather합니다.

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

전략	설명
`ColwiseParallel`	가중치를 출력 특징 차원에서 샤딩하고(2D `nn.Linear` 가중치의 경우 `Shard(0)`, `nn.Embedding`의 경우 `Shard(1)`), 1D 편향을 샤딩합니다. 입력을 `Replicate()`로 재배포한 뒤 출력을 `output_layouts`에 배치합니다.
`RowwiseParallel`	가중치를 입력 특징 차원에서 샤딩하고(`Shard(-1)`, `nn.Embedding`의 경우 `Shard(0)`), 축소 후 추가되는 편향을 복제합니다. 각 랭크의 forward는 합에서 자신의 몫을 보유한 `Partial()` 출력을 만들며, `transform_output_post_forward`이 이를 `output_layouts`로 재배포합니다. `Replicate()`로 축소하면 all-reduce가 발생하고, `Shard(1)`로 축소하면 reduce-scatter가 발생합니다.
`PackedColwiseParallel`	융합된 가중치를 위한 `ColwiseParallel` 변형입니다. 예를 들어 `up_proj`와 `gate_proj`가 `gate_up_proj`로 패킹된 경우입니다. 가중치가 2개가 아닌 여러 레이어를 패킹할 때는 `split_factor`를 사용하세요.
`PackedRowwiseParallel`	마지막 차원을 따라 패킹된 가중치를 위한 행 방향 대응물입니다. 1D 파라미터를 복제합니다.
`SequenceParallel`	모듈의 파라미터를 복제하고 입력을 `sequence_dim`(기본값 `1`)에서 샤딩합니다. `LayerNorm`과 `RMSNorm`처럼 토큰별로 동작하는 norm에 사용됩니다.
`ReplicatedWithGradAllReduce`	파라미터를 복제하지만 그 그래디언트를 all-reduce합니다. 열 방향과 행 방향 레이어 사이에 있고 샤딩된 축을 따라 정규화하는 norm에 필요합니다. 각 랭크는 자신의 헤드만 보기 때문이에요.
`AllReduceParallel`	모듈의 `Partial()` forward 출력을 `Replicate()`로 all-reduce합니다. 계산이 부분 합으로 끝나는 모듈의 동기화 지점으로 사용하세요.
`MlaKvAProjParallel`	DeepSeek-V2 스타일 MLA 어텐션의 `kv_a_proj_with_mqa` 출력을 분할하고 RoPE 절반의 그래디언트를 all-reduce합니다. 이는 `kv_b_proj`를 우회하는데, 그대로 두면 부분 그래디언트가 유지되기 때문입니다. 모델 설정에 `qk_rope_head_dim`이 필요해요.
`MoEParamShard`	주어진 배치에서 MoE 전문가 가중치를 샤딩합니다. `grouped_gemm` 이름을 뒷받침하며, `shards_expert_dim=True`는 `module.num_experts`를 랭크별 전문가 수로도 다시 씁니다.
`EpRouterParallel`	비로컬 전문가에 대한 라우터 점수를 마스킹하고 전역 전문가 ID를 로컬 ID로 다시 매핑해서, 각 랭크가 소유한 전문가만 실행하게 합니다. `num_experts`가 메시 크기로 나누어떨어져야 해요.
`RouterParallelMegaMoe`	DeepGEMM Mega MoE를 위한 라우터 변형으로, 라우터 출력을 그대로 두기를 원하는 커널 내부에서 전문가를 디스패치합니다.
`MoeExpertsParallel`	텐서 병렬 MoE 전문가. 전문가 출력 forward를 all-reduce하고 숨은 상태와 라우팅 가중치에 대한 backward all-reduce를 추가합니다.
`MoeTensorParalellMegaMoeExperts`	DeepGEMM Mega MoE를 위한 추론 전용 전문가 레이어. 그래디언트 동기화를 건너뛰고 프로세스 그룹을 모듈에 전달해서 커널이 첫 forward에서 공유 버퍼를 설정할 수 있게 합니다.
`MoeIdentityParallel`	제로 또는 아이덴티티 전문가의 입력을 메시 크기로 미리 나눠서, `moe_tp_experts`가 하위에서 적용하는 all-reduce를 상쇄합니다.

패킹 전략 (Packed strategies)

가중치 패킹은 여러 선형 레이어를 단일의 더 큰 레이어로 결합합니다. PackedColwiseParallel과 PackedRowwiseParallel 전략은 패킹된 가중치를 올바르게 샤딩합니다. 기본 ColwiseParallel이나 RowwiseParallel 전략은 패킹된 가중치를 잘못 샤딩해요.

아래 예시는 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는 그 차원을 연속 블록으로 나누기 때문에, 랭크는 각각의 슬라이스 대신 gate의 꼬리와 up의 머리를 받게 됩니다. 패킹 전략은 split_factor(기본값 2)와 함께 _StridedShard를 사용해 분할을 인터리브하여, 모든 랭크가 두 절반의 일치하는 슬라이스를 받게 하므로 샤딩 후에도 chunk가 여전히 맞아떨어집니다.

커스텀 분할 전략 (Custom partitioning strategies)

커스텀 분할 전략을 만들려면 distributed/tensor_parallel.py의 TensorParallelLayer를 상속하세요. 전략에 필요한 훅만 재정의하면 됩니다. 모든 훅에는 no-op 기본값이 있기 때문이에요.

훅	용도
`validate_param`	가중치가 로드되기 전에, 전략이 샤딩할 수 없는 파라미터를 거부합니다.
`shard_param`	파라미터 하나를 `DTensor` 플레이스홀더로 교체해 로더가 어떤 샤드가 이 랭크에 속하는지 알게 합니다.
`transform_inputs_pre_forward`	모듈의 입력을 forward가 기대하는 레이아웃으로 재배포합니다.
`context_around_forward`	forward를 컨텍스트 매니저로 감쌉니다. 예를 들어 로컬 텐서를 커널에 노출할 때 사용합니다.
`transform_output_post_forward`	모듈의 출력을 재배포하거나 축소합니다.
`should_use_local_tensors`	이 모듈의 forward가 `DTensor`가 아닌 평범한 텐서가 필요함을 보고합니다.
`install_forward`	`module.forward`를 완전히 교체합니다. 위 훅만으로 충분하지 않을 때만 재정의하세요. `ReplicatedWithGradAllReduce`가 backward 훅을 등록할 때처럼요.

아래 예시는 ColwiseParallel의 축소 버전을 단계별로 보여줍니다.

TensorParallelLayer를 상속하고 전략이 사용하는 배치를 저장합니다. 기본 클래스는 __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

파라미터 하나를 DTensor 플레이스홀더로 감싸도록 shard_param을 구현합니다. meta 텐서에서 실행되므로 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,
    )

입력·출력 변환을 구현합니다. 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 왕복을 건너뜁니다. 전략에 같은 처리가 필요하다면 복사하기 전에 소스를 읽어 보세요.

tp_plan이 이름을 붙일 수 있도록 전략을 등록합니다. 등록은 전략 인스턴스를 받으므로 __init__에 넘긴 배치가 이름과 함께 저장됩니다. 아래의 ()에 주목하세요.

import torch

from transformers import AutoModelForCausalLM, DistributedConfig
from transformers.distributed.tensor_parallel import ParallelInterface

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

벤치마크 (Benchmarks)

텐서 병렬화는 특히 큰 배치 크기나 긴 시퀀스에서 추론을 크게 가속합니다.

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

설계 구현 (Design implementation)

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(분산 텐서)는 일반적인 텐서 연산 위에서 분산 로직을 처리합니다. 텐서 병렬화 아래의 모델 가중치는 DTensor로 저장되며, 그래서 전략이 통신을 명시적인 컬렉티브 대신 레이아웃 변경으로 설명할 수 있는 것입니다.

placements 속성은 텐서가 DeviceMesh의 기기들 사이에 어떻게 배치되는지 PyTorch에 알려줍니다. 다음 값을 받아들입니다.

Shard(dimension)는 DTensor를 만들어졌을 때의 DeviceMesh 위에서 주어진 차원으로 분할합니다. 열 방향 분할은 가중치의 출력 특징 차원과 편향의 유일한 차원을 샤딩합니다.

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

행 방향 분할은 대신 입력 특징 차원을 샤딩하고 편향을 복제합니다. 편향은 모든 랭크가 아니라 축소 후에 한 번 추가되기 때문입니다.

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를 발생시키는 것입니다.

더 알아보기 (Learn more)

텐서 병렬화에 대한 Ultra-Scale Playbook 섹션에 더 많은 세부 정보가 있습니다.
mixture-of-experts(MoE) 모델을 사용한다면 expert parallelism 가이드를 확인하세요. 이런 모델은 텐서 병렬화와 전문가 병렬화를 지원합니다.
Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand 블로그 포스트를 읽으면 텐서 병렬화를 빠르게 파악하고 열·행 병렬 설정이 어떻게 다른지 배울 수 있어요.
훈련 환경에서 사용하는 방법은 Tensor parallelism 훈련 가이드를 보세요.