Transformers documentation
분산 추론
분산 추론
모델이 단일 GPU에 올라가지 않는 경우, 텐서 병렬 처리를 사용한 분산 추론이 도움이 될 수 있습니다. 텐서 병렬화는 모델을 여러 가속기(CUDA GPU, Intel XPU 등)에 분할하여 행렬 곱셈과 같은 계산을 병렬화합니다. 이를 통해 더 큰 모델을 메모리에 올릴 수 있으며, 각 가속기가 텐서의 일부를 처리하므로 추론 속도가 향상됩니다.
그러나 텐서 병렬화는 통신 오버헤드를 발생시키므로, 빠른 노드 내 통신을 활용할 수 있는 다중 가속기 환경에서 사용하는 것이 가장 효과적입니다. 다중 노드 학습 환경에서는 사용 사례에 따라 파이프라인 병렬화나 데이터 병렬화를 사용하는 것이 더 효율적일 수 있습니다.
텐서 병렬화에 대해 더 자세히 알아보려면 Ultra-Scale Playbook의 텐서 병렬화 섹션을 참조하세요.
아래 목록에서 텐서 병렬 처리를 기본적으로 지원하는 모델을 확인할 수 있습니다. 새로운 모델에 대한 지원을 추가하려면 GitHub 이슈나 풀 리퀘스트를 열어주세요.
지원되는 모델 보기
이 가이드는 Transformers에서 다양한 분할 전략을 사용하여 텐서 병렬화를 활성화하는 방법을 설명합니다.
모델 분할
Transformers는 tp_plan매개변수를 활용할 수 있는 모델에 대해 텐서 병렬 처리를 지원합니다. 모델 분할 방식은 두 가지가 있습니다.
auto텐서 병렬화 계획은 사전 정의된 구성을 기반으로 모델(위에 언급된 지원 모델)을 자동으로 분할합니다.- 사용자 지정 분할 계획을 직접 정의하여 [~PreTrainedModel.from_pretrained] 메소드의
tp_plan매개변수로 전달할 수 있습니다.
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # 모든 가능한 전략을 시각화하기에 더 좋음
model_id = "meta-llama/Meta-Llama-3-8B-Instruct" # 적은 수의 GPU에 더 좋음
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, tp_plan="auto")
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)
# 분산 실행
outputs = model(inputs)위의 추론 스크립트를 GPU당 4개 프로세스로 torchrun에서 실행하세요.
torchrun --nproc-per-node 4 demo.py
분할 전략
모든 분할 전략은 문자열을 전략 구현에 매핑하는 ParallelInterface 클래스에서 정의됩니다. 모든 전략은 from_pretrained()의 tp_plan을 통해 설정되므로 이 클래스와 직접 상호 작용할 필요는 없지만, 어떤 전략을 사용할 수 있는지 확인할 때 유용합니다.
class ParallelInterface(MutableMapping):
"""
허용된 어텐션 함수를 추적하는 딕셔너리 같은 객체입니다. `register()` 호출로 새로운 어텐션 함수를 쉽게 추가할 수 있습니다.
모델이 기존 어텐션 함수(예: `sdpa`)를 로컬에서 덮어쓰려면 `modeling_<model>.py` 내부에서 이 클래스의 새 인스턴스를 선언하고
해당 인스턴스에서 선언해야 합니다.
"""
_global_mapping = {
"colwise": ColwiseParallel(),
"rowwise": RowwiseParallel(),
"colwise_rep": ColwiseParallel(output_layouts=Replicate()),
"rowwise_rep": RowwiseParallel(input_layouts=Replicate()),
"local_colwise": ColwiseParallel(use_dtensor=False),
"local_rowwise": RowwiseParallel(use_dtensor=False),
"local": IsolatedParallel(),
"gather": GatherParallel(),
"local_packed_rowwise": PackedRowwiseParallel(use_dtensor=False),
"sequence_parallel": SequenceParallel(),
"replicate": ReplicateParallel(),
}각 전략에 대해 자세히 알아보려면 아래 표를 참고하세요.
| 전략 | 설명 |
|---|---|
ColwiseParallel | 가중치와 편향의 열 방향 분할. |
RowwiseParallel | 가중치와 편향의 행 방향 분할. nn.Embedding 모듈 분할도 지원. |
SequenceParallel | LayerNorm과 Dropout 레이어를 지원하는 시퀀스 병렬 구현. RMSNorm의 Python 구현도 지원. |
PackedColwiseParallel | 패킹된 가중치를 지원하는 ColwiseParallel의 변형(예: up_proj와 gate_proj를 함께 패킹). 자세한 내용은 코드를 참조하세요. |
PackedRowwiseParallel | 패킹된 가중치를 지원하는 RowwiseParallel의 변형(코드 참조). |
GatherParallel | 기기 간 모듈의 출력을 수집. |
IsolatedParallel | Mixture-of-Experts(MoE) 레이어의 전문가에 사용되어 다른 기기로부터 모듈을 격리. |
ReplicateParallel | 부분적으로 분할된 모델로 인해 torch.distributed API가 중단되는 것을 방지하기 위해 모든 기기에 모듈을 복제. |
패킹된 전략
가중치 패킹은 여러 선형 레이어를 하나의 더 큰 레이어로 합치는 기법입니다. 패킹된 전략인 PackedColwiseParallel과 PackedRowwiseParallel은 패킹된 가중치를 분할하는 데 사용됩니다. 기본적인 ColwiseParallel이나 RowwiseParallel은 패킹된 가중치를 올바르게 분할하지 못합니다.
아래 예시는 up_proj와 gate_proj를 단일 gate_up_proj 모듈로 패킹하고 gate_up_proj를 분할하기 위해 PackedRowwiseParallel 전략이 필요합니다.
class Llama4TextExperts(nn.Module):
...
self.gate_up_proj = nn.Parameter(torch.empty(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) # gate_up_proj 모듈의 출력 계산
gate, up = gate_up.chunk(2, dim=-1) # 출력을 gate와 up으로 분할Packed*를 사용해야 하는 이유에 대한 시각적 표현은 이 주석을 참고하세요.
로컬 전략
로컬 전략(local_colwise, local_rowwise, local_packed_rowwise)은 torch.chunk와 같은 일부 연산에서 지원되지 않기 때문에 DTensor를 사용하지 않습니다. 대신 로컬 전략은 기본 torch.Tensor를 사용하고 일부 분산 로직을 수동으로 수행합니다.
사용자 정의 분할 전략
사용자 정의 분할 전략은 TensorParallelLayer를 상속하고 partition_tensor, _prepare_input_fn, _prepare_output_fn을 구현해야 합니다.
그런 다음 tp_plan에서 해당 전략을 지정했을 때 디스패칭 로직이 찾을 수 있도록 ParallelInterface 매핑에 등록해야 합니다.
아래 예시는 이 워크플로우로 ColwiseParallel을 구현하는 방법을 보여줍니다.
TensorParallelLayer를 상속합니다.__init__메소드에서 입력 및 출력 텐서가 기기에 어떻게 배치되어야 하는지 설명하는input_layouts과output_layouts을 정의합니다.desired_input_layouts속성은 입력이 기기에 어떻게 배치되어야만 하는지를 명시하는 데 사용됩니다.class ColwiseParallel(TensorParallelLayer): def __init__( self, *, input_layouts: Optional[Placement] = None, # 이전 레이어에서 오는 입력 레이아웃 output_layouts: Optional[Placement] = None, # 달성하고자 하는 출력 레이아웃 use_local_output: bool = True, # 로컬 출력 사용 여부 use_dtensor=True, # DTensor 사용 여부 ): self.input_layouts = (input_layouts or Replicate(),) # 이전 레이어에서 오는 입력 분할 self.output_layouts = (output_layouts or Shard(-1),) # 원하는 출력 분할 self.desired_input_layouts = (Replicate(),) # 원하는 입력 분할, 입력은 GPU 간에 복제되어야 함 self.use_local_output = use_local_output self.use_dtensor = use_dtensorpartition_tensor,_prepare_input_fn,_prepare_output_fn메서드를 구현합니다.partition_tensor메소드는 텐서를 분할하고 분할된 텐서로empty_param을 채웁니다. 유틸리티 함수get_tensor_shard를 사용하여 주어진 랭크에 대한 원본 매개변수의 올바른 분할을 얻고, 패킹된 가중치에 대해서는get_packed_weights를 사용하세요.def partition_tensor( self, param, # 매개변수의 전체 텐서 empty_param, # 매개변수의 빈 텐서, 분할된 텐서로 채워짐 param_type, # 매개변수 유형, `bias` 또는 `weight` param_casting_dtype, # 매개변수를 캐스팅할 유형 to_contiguous, # 텐서를 연속적인 메모리 레이아웃으로 변환할지 여부 rank, # 현재 기기의 랭크 device_mesh, # 기기 메시 ) -> nn.Parameter: # 분할된 매개변수 반환 ..._prepare_input_fn과_prepare_output_fn메소드는 사전 포워드 및 포워드 훅에서 사용됩니다.__init__에서 지정된 대로 입력과 출력을 원하는 레이아웃으로 재분배합니다.def _prepare_input_fn(input_layouts, desired_input_layouts, mod, inputs, device_mesh): ... # 사용자 정의 로직 수행, DTensor로 캐스팅 등. ... return inputs.redistribute(placements=desired_input_layouts, device_mesh=device_mesh) def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh): ... # 사용자 정의 로직 수행, DTensor로 캐스팅 등. ... return outputs.redistribute(placements=output_layouts, device_mesh=device_mesh)tp_plan과 함께 사용할 수 있도록 전략을ParallelInterface에 등록합니다.from transformers.integrations.tensor_parallel import ParallelInterface ParallelInterface.register_strategy("colwise_custom", ColwiseParallel) tp_plan = { "model.layers.*.self_attn.q_proj": "colwise_custom", ... } model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, tp_plan=tp_plan)
벤치마크
텐서 병렬화는 특히 큰 배치 크기나 긴 시퀀스를 가진 입력에 대한 추론 속도를 크게 향상시킬 수 있습니다.
시퀀스 길이가 512인 Llama에서 단일 포워드 패스에 대한 예상 속도 향상 수치는 아래 차트를 참조하세요.
설계 구현
Transformers 텐서 병렬화 구현은 프레임워크에 구애받지 않지만, 구체적인 구현을 위해서는 DeviceMesh와 torch.distributed의 DTensor에 의존하여 간단하고 확장 가능한 인터페이스를 제공합니다.
DeviceMesh
DeviceMesh를 함께 통신하는 기기들의 다차원 그리드로 상상해보세요. 병렬 처리 전략마다 각기 다른 통신 패턴이 필요하므로, 여러 하위 메시를 가진 DeviceMesh를 만들 수 있습니다.
from torch.distributed.device_mesh import init_device_mesh
# 4개 GPU의 1D 메시 생성
device_mesh = init_device_mesh("cuda", (4,), mesh_dim_names=["tp"])torch.distributed에서 정의된 대부분의 병렬화 전략은 메시 자체나 하위 메시에 적용할 수 있으며, 자동으로 통신 패턴을 처리합니다.
DTensor
DTensor(분산 텐서)는 일반적인 텐서 연산 위에 분산 로직을 처리하는 텐서 하위 클래스입니다. 텐서 병렬화의 대부분의 모델 가중치는 DTensor 형태로 저장됩니다.
DTensor의 가장 중요한 부분은 placement 속성입니다. 이는 PyTorch에게 텐서가 DeviceMesh의 기기에 어떻게 배치되는지 알려주기 때문입니다. placement 속성은 다음 값을 가질 수 있습니다.
Shard(dimension)-DTensor가 구성된DeviceMesh에서 주어진 차원에 걸쳐 어떻게 분할되는지 나타냅니다. 아래 예시는 열 방향 분할을 위해 다양한 차원에 걸쳐 가중치를 분할하는 방법을 보여줍니다.weight = ... weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # 첫 번째(열 방향) 차원에 걸쳐 분할 bias = ... bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # 유일한 차원에 걸쳐 분할이 예시는 행 방향 분할을 위해 여러 차원에 걸쳐 가중치를 분할하는 방법을 보여줍니다.
weight = ... weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(1)]) # 두 번째(행 방향) 차원에 걸쳐 분할 bias = ... bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # 모든 GPU에 편향 복제Replicate()-DTensor가DeviceMesh에 걸쳐 복제됨을 나타냅니다. 각 기기에 텐서의 전체 사본만 생성합니다.bias = ... bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # 모든 GPU에 편향 복제Partial()- 텐서가 감소 연산을 기다리고 있는 상태임을 나타냅니다 (일반적으로 Transformers에서의 사용 사례와는 직접적인 관련이 적습니다).