# 추론을 위한 텐서 병렬성

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

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

```py
from transformers import AutoModelForCausalLM

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

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

## 모델 파티셔닝

[DistributedConfig](/docs/transformers/main/ko/expert_parallelism#transformers.DistributedConfig)에서 `tp_size`로 텐서 병렬 장치 수를 설정합니다.

- 모델의 사전 정의된 플랜을 사용하려면 `DistributedConfig(tp_size=N)`를 설정하세요.
- 수동으로 `tp_plan`을 정의하고 `tp_size`와 함께 [DistributedConfig](/docs/transformers/main/ko/expert_parallelism#transformers.DistributedConfig)에 전달합니다.

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

```py
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](https://pytorch.org/docs/stable/elastic/run.html)로 추론 스크립트를 실행합니다. GPU당 프로세스를 하나씩 사용합니다.

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

`tp_plan`의 각 레이어에 대해 텐서 병렬 계획을 정의하고 [DistributedConfig](/docs/transformers/main/ko/expert_parallelism#transformers.DistributedConfig)를 통해 전달합니다. 아래 예제는 column 및 row 분할을 사용합니다. 다른 지원되는 전략은 [Partitioning strategies](#partitioning-strategies) 섹션을 참조하세요.

수동 파티셔닝은 모델 아키텍처와 전략 간 상호작용에 대한 깊은 이해가 필요합니다. 파티셔닝 선택이 부적절하면 실패하거나 잘못된 결과를 생성하는 느린 모델이 만들어집니다. [Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism)에서는 파티셔닝 전략을 자세히 설명합니다.

Keys는 모듈 또는 파라미터 이름이며, `*`은 레이어 인덱스를 나타냅니다. 인식되지 않는 전략 이름은 지원되는 이름을 나열한 `ValueError`를 발생시키며, Transformers는 어떤 대상에도 일치하지 않는 plan 규칙과 어떤 규칙에도 포함되지 않는 파라미터에 대해 경고를 기록합니다.

```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](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/parallel_utils.py)에 정의된 `ParallelInterface`로, 플랜에서 사용할 수 있는 각 전략 이름을 구성된 전략 인스턴스에 매핑합니다. 모델을 샤딩하는 데 직접 상호작용하지는 않지만, 사용 가능한 이름의 공식 목록입니다. [distributed/tensor_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py)는 텐서 병렬 전략을 여기에 등록합니다:

```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](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/expert_parallel.py)에 있으며, 동일한 훅을 가진 `ExpertParallelLayer`를 하위 클래스로 상속하고 동일한 `ALL_PARALLEL_STYLES`에 등록됩니다. 가장 자주 사용할 두 가지인 `ColwiseParallel`와 `RowwiseParallel`은 `input_layouts`와 `output_layouts` 배치(placement)를 받아, 하나의 클래스가 여러 plan 이름을 커버할 수 있습니다. `colwise`는 출력을 마지막 차원별로 샤드(shard) 상태로 남겨 두어 뒤따르는 `rowwise` 레이어가 소비하도록 하고, `colwise_gather_output`는 전체 텐서로 다시 all-gather합니다.

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

| 전략 | 설명 |
|---|---|
| `ColwiseParallel` | 선형 가중치를 출력 특징 차원(`Shard(0)`은 2차원 `nn.Linear` 가중치에 해당) 기준으로 샤딩하고, 1차원 바이어를 샤딩합니다. 입력을 `Replicate()`로 재분배한 후, 출력을 `output_layouts`에 배치합니다. |
| `RowwiseParallel` | 입력 특징 차원(`Shard(-1)`)에 선형 가중치를 샤딩하고, 리덕션 이후에 추가되는 바이어스는 복제합니다. 각 랭크의 포워드가 합계의 해당 랭크 몫을 담은 `Partial()` 출력을 생성하면, `transform_output_post_forward`가 이를 `output_layouts`로 재분배합니다. `Replicate()`로 리듀스하면 올리듀스가, `Shard(1)`로 리듀스하면 리듀스-스캐터가 실행됩니다. |
| `EmbeddingParallel` | 토큰 ID가 복제된 `nn.Embedding`을 샤딩합니다. `shard_dim=0`(`embedding_rowwise`, 일반적인 선택)은 어휘를 샤딩하고 각 랭크의 부분 조회 결과를 전체 출력으로 올리드(all-reduce)합니다. `shard_dim=1`(`embedding_colwise`)은 숨겨진 차원을 샤딩하고 각 랭크의 출력 슬라이스를 올가더(all-gather)합니다. |
| `PackedColwiseParallel` | 융합된 가중치를 위한 `ColwiseParallel`의 변형입니다. 예를 들어 `up_proj`와 `gate_proj`가 `gate_up_proj`로 묶인 경우입니다. 가중치가 두 개가 아닌 다른 수의 레이어를 포함할 때는 `split_factor`를 사용하세요. |
| `PackedRowwiseParallel` | 최종 차원 방향으로 가중치가 패킹된 행 단위 대응 함수입니다. 1D 매개변수를 복제합니다. |
| `SequenceParallel` | 모듈의 매개변수를 복제하고 입력을 `sequence_dim`(기본값은 `1`)을 기준으로 샤드합니다. `LayerNorm` 및 `RMSNorm`과 같이 토큰별로 동작하는 정규화에 사용됩니다. |
| `ReplicatedWithGradAllReduce` | Replicates a parameter but all-reduces its gradient. Needed for norms that sit between a column-wise and a row-wise layer and normalize along a sharded axis, where each rank only sees its own heads. |
| `AllReduceParallel` | 모듈의 `Partial()` forward 출력을 `Replicate()`로 all-reduce합니다. 계산이 부분 합으로 끝나는 모듈의 동기화 지점으로 사용하세요. |
| `MlaKvAProjParallel` | DeepSeek-V2 스타일 MLA attention의 `kv_a_proj_with_mqa` 출력을 분리하고, RoPE 절반의 그래디언트를 all-reduce합니다. 이 그래디언트는 `kv_b_proj`를 우회하므로, 처리하지 않으면 부분 그래디언트가 유지됩니다. 모델 설정에 `qk_rope_head_dim`이 필요합니다. |
| `MoeExpertsParallel` | Tensor parallel MoE experts. expert output의 forward를 all-reduce하고, hidden states와 routing weights에 대한 backward all-reduce를 추가합니다. |
| `MoeIdentityParallel` | `moe_tp_experts`가 하류에서 적용하는 all-reduce를 상쇄하기 위해, zero 또는 identity expert의 입력을 mesh size로 미리 나눕니다. |
| `LocalParamsParallel` | 모듈의 forward를 자체 샤드된 파라미터를 로컬 텐서로 사용하여 실행합니다. 자체 플랜 엔트리가 샤드하는 attention sinks와 같은 순수한 `nn.Parameter`의 소유자에 사용하세요. |
| `MoEParamShard` | 주어진 placement에 따라 MoE expert 가중치를 샤딩합니다. `grouped_gemm` 이름을 뒷받침하며, `shards_expert_dim=True`인 경우 `module.num_experts`를 랭크별 expert 수로 다시 작성합니다. |
| `EpRouterParallel` | 로컬이 아닌 전문가의 router scores를 mask하고, 전역 전문가 ID를 로컬 ID로 재매핑하여 각 rank가 자신이 소유한 전문가만 실행하도록 합니다. `num_experts`이 mesh size로 나누어떨어져야 합니다. |
| `RouterParallelMegaMoe` | DeepGEMM Mega MoE의 Router 변형으로, 커널 내부에서 전문가를 디스패치하며 라우터 출력이 그대로 유지되기를 원하는 경우를 위해 제공됩니다. |
| `MoeTensorParalellMegaMoeExperts` | DeepGEMM Mega MoE의 추론 전용 전문가(experts) 레이어입니다. 그라디언트 동기화를 건너뛰고 프로세스 그룹을 모듈로 전달하여, 첫 번째 포워드 패스에서 커널이 공유 버퍼를 설정할 수 있도록 합니다. |

### 패킹 전략

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` 전략이 필요합니다.

```python
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` 모듈의 출력을 계산합니다.

```python
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](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py)에서 상속하여 커스텀 파티셔닝 전략을 만듭니다. 모든 훅에는 no-op 기본값이 있으므로, 전략에 필요한 훅만 오버라이드하세요.

| Hook | 목적 |
|---|---|
| `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`가 백워드 훅을 등록하는 방식이 그 예입니다. |

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

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

    ```python
    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`을 전달합니다.

    ```python
    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 주위에 이를 호출하므로, 텐서를 레이아웃 간에 이동하는 것만 수행하면 됩니다. 열 단위 파티셔닝은 복제된 입력을 기대하고 마지막 차원별로 샤딩된 출력을 생성합니다.

    ```python
    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가 해당 이름과 함께 저장됩니다. 아래 `()`에 주의하세요.

    ```python
    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](./model_doc/llama)에서 단일 forward pass에 대한 예상 속도 향상을 보여줍니다.

    

## 설계 구현

Transformers는 프레임워크에 종속되지 않는 방식으로 텐서 병렬성을 구현합니다. [torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html)의 [DeviceMesh](https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html)와 [DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html)에 의존하여 단순하고 확장 가능한 인터페이스를 제공합니다.

### DeviceMesh

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

```python
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`를 주어진 차원을 기준으로 분할합니다. 열 방향 파티셔닝은 가중치의 출력 특징 차원과 바이어스의 유일한 차원을 샤딩합니다.

    ```python
    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 이후에 한 번만 추가되기 때문입니다.

    ```python
    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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism) 섹션에서 확인하세요.

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

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

- 학습 환경에서 사용하는 방법을 배우려면 [Tensor parallelism](./tensor_parallelism) 학습 가이드를 참조하세요.