Buckets:

hf-doc-build/doc / diffusers /main /ko /optimization /tensor_parallelism.md
mishig's picture
|
download
raw
11.8 kB
# Tensor 병렬성
[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism)은 모델의 가중치 행렬을 여러 디바이스에 걸쳐 분할합니다. 각 디바이스는 각 레이어의 열 방향(`"colwise"`) 또는 행 방향(`"rowwise"`) 슬라이스를 보유하며 부분 결과를 계산합니다. 레이어 경계에서의 `AllReduce` 또는 `AllGather` 연산은 부분 결과를 결합하여 전체 출력을 생성합니다. [context parallelism](./context_parallelism)과 달리, 이는 디바이스당 *weight* 메모리를 줄여주므로 모델이 단일 디바이스에 맞지 않을 때 유용합니다. 다른 전략과의 비교를 위해 [Distributed inference](../training/distributed_inference)를 참조하세요.
## Load a sharded model
Pass a [TensorParallelConfig](/docs/diffusers/main/ko/api/parallel#diffusers.TensorParallelConfig) to the `parallel_config` argument of the model's [from_pretrained()](/docs/diffusers/main/ko/api/models/overview#diffusers.ModelMixin.from_pretrained). `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model's class must define a `_tp_plan`, which maps module names to a sharding style. To check whether a model supports tensor parallelism, look for this attribute on its class, such as [Flux2Transformer2DModel](/docs/diffusers/main/ko/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).
다른 모델에 플랜을 추가하려면 [모델에 텐서 병렬성 추가](#add-tensor-parallelism-to-a-model)를 참조하세요.
이 방식으로 로드하면 체크포인트가 읽히는 동안 샤드됩니다. 각 랭크는 각 가중치의 자기 자신의 슬라이스만 읽고 이를 해당 GPU에 직접 배치하므로, 어떤 랭크도 전체 크기의 가중치를 보유하지 않습니다. 랭크당 메모리는 `tp_degree`가 증가함에 따라 감소합니다.
```py
import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig
def main():
dist.init_process_group(backend="nccl")
rank, world_size = dist.get_rank(), dist.get_world_size()
# torchrun sets LOCAL_RANK to this process's GPU index on its node
local_rank = int(os.environ["LOCAL_RANK"])
device = torch.device(f"cuda:{local_rank}")
torch.cuda.set_device(device)
# Each rank reads only its own shard of every planned weight, straight onto its own GPU.
transformer = Flux2Transformer2DModel.from_pretrained(
"black-forest-labs/FLUX.2-dev",
subfolder="transformer",
dtype=torch.bfloat16,
parallel_config=TensorParallelConfig(tp_degree=world_size),
)
pipeline = DiffusionPipeline.from_pretrained(
"black-forest-labs/FLUX.2-dev", transformer=transformer, dtype=torch.bfloat16
)
# The transformer is already on its device; move the remaining components individually. Do not call
# `pipeline.to(device)` — that would move every rank's shards onto the same device.
pipeline.text_encoder.to(device)
pipeline.vae.to(device)
generator = torch.Generator().manual_seed(42)
image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
if rank == 0:
image.save("output.png")
dist.destroy_process_group()
if __name__ == "__main__":
main()
```
스크립트를 `tensor_parallel_flux2.py`로 저장하고 `torchrun`으로 실행하세요.
```shell
torchrun --nproc-per-node 4 tensor_parallel_flux2.py
```
스크립트는 `tp_degree`를 `world_size`로 설정하므로, `--nproc-per-node 4`는 transformer를 4개의 GPU에 걸쳐 샤딩합니다.
transformer만 샤딩됩니다. 각 랭크는 텍스트 인코더의 전체 사본을 보유하므로, 각 GPU는 텍스트 인코더와 transformer의 해당 랭크 몫을 위한 충분한 메모리가 필요합니다.
## 로딩 성능
`parallel_config`를 사용하여 로드하는 방식은 전체 모델을 로드한 후 `enable_parallelism()`을 호출하는 방식보다 빠르고, 랭크당 CPU 메모리 사용량이 적습니다. 아래 표는 4x A10G (23GB) 환경에서 [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (transformer만, 32B 파라미터, bf16)에 대한 결과를 보여줍니다.
| tp_degree | method | 로드 시간 | peak CPU/rank |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB |
## 한계
Tensor parallelism은 일부 로딩 옵션 및 기능과 함께 작동하지 않으며, 샤드된 모델을 저장할 수 없습니다.
- `parallel_config`을 사용한 텐서 병렬 로딩은 `device_map`, `quantization_config`, `low_cpu_mem_usage=False`, `use_flashpack=True`, 또는 safetensors가 아닌 가중지와 함께 작동하지 않습니다. 각각의 경우 전체 체크포인트를 조용히 로드하는 대신 오류가 발생합니다.
- tensor parallelism을 양자화, 오프로딩 또는 LoRA 어댑터와 결합하는 경우에도 모델이 어떻게 샤딩되든 오류가 발생합니다.
- 메모리에 이미 로드된 모델을 샤딩하려면 동일한 구성으로 `enable_parallelism()`을 호출합니다. 이 방법은 먼저 전체 모델을 로드한 후 샤딩을 수행하므로, 모든 랭크에 전체 체크포인트를 위한 메모리가 필요합니다. `parallel_config`를 사용하여 로드하면 이러한 비용을 피할 수 있습니다.
- 텐서 병렬 모델의 저장은 [save_pretrained()](/docs/diffusers/main/ko/api/models/overview#diffusers.ModelMixin.save_pretrained)에서 오류를 발생시킵니다. 샤딩하기 전에 모델을 저장하세요.
## 모델에 텐서 병렬성 추가
이 섹션은 Diffusers 모델에 텐서 병렬성(tensor parallelism) 지원을 추가하려는 기여자를 위한 것입니다.
Tensor parallelism은 `_tp_plan`이라는 클래스 속성을 정의한 모델에서만 작동합니다. 이 속성은 모듈 이름을 샤딩 스타일에 매핑합니다. 이 속성을 작성하는 것은 대부분 *확장*하는 투영과 *축소*하는 투영을 짝짓는 것을 의미합니다.
각 키에는 `*`가 최대 하나만 포함될 수 있으며, `*` 앞부분은 [`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html)로 해석되어야 하므로 하나의 항목이 모든 블록을 커버합니다. `*`가 없는 키는 모델 자체에 적용됩니다. 경로는 모델 기준 상대 경로입니다.
### Colwise와 rowwise
| 스타일 | Shards | 각 랭크 | 사용 |
|---|---|---|---|
| `"colwise"` | 출력 특징(`weight` 차원 0) | 출력의 슬라이스를 계산합니다 | `to_q`, `to_k`, `to_v`, FFN in-projection |
| `"rowwise"` | 입력 특징(`weight` 차원 1) | 부분 합을 계산합니다 | `to_out.0`, FFN 출력 투사 |
| `"rowwise_split_input"` | 입력 특징(`weight` 차원 1) | 복제된 입력의 슬라이스에서 부분 합을 계산합니다 | 입력이 복제된 독립적인 투영, 예를 들어 타임스텝 임베딩을 읽는 AdaLN 변조 투영 |
`"colwise"`와 `"rowwise"`는 항상 이 순서대로 짝을 이루어야 합니다. `"colwise"` 투영은 출력을 샤드 상태로 유지하며, 바로 다음 `"rowwise"` 투영은 해당 샤드를 직접 소비하고, 블록 경계에서 한 번의 `AllReduce`만 수행하면 결과를 재구성할 수 있습니다. 이 쌍을 다른 방식으로 샤딩하면 중간에 gather가 강제되어 통신량이 훨씬 커집니다.
Attention에서 열 방향 Q, K, V 투영은 각 rank에 head의 부분 집합을 할당하므로 `tp_degree`은 head 수를 나누어야 합니다. Encoder-stream 복제 항목(`add_q_proj`, `to_add_out`, `ff_context`)은 image-stream의 대응 항목과 동일한 패턴을 따릅니다.
### Fused projections
`Linear` 하나가 샤딩 대상 차원을 따라 여러 논리 텐서를 패킹하는 경우, 단순한 `"colwise"`/`"rowwise"` 슬라이싱은 연결된 부분을 가로질러 잘라내어 조각들이 어긋나게 됩니다. 대신 각 패킹된 블록을 독립적으로 샤딩하는 `PackedColwiseParallel`/`PackedRowwiseParallel`을 사용하세요.
```py
# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
```
`blocks`는 packed 차원을 나누는 합을 가지는 비율 정수 목록입니다. 예를 들어, 동일한 절반으로 구성된 SwiGLU 게이트 및 업 프로젝션에는 `[1, 1]`을, `mlp_ratio=3`인 융합된 Q, K, V, 게이트 및 업 프로젝션에는 `[1, 1, 1, 3, 3]`을 사용합니다.
```py
"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),
```
`blocks`가 config에서만 알려진 경우, `blocks`를 생략하고 대신 `__init__` 중에 `Linear`에 절대 크기를 `_tp_packed_col_blocks` 또는 `_tp_packed_row_blocks`로 저장합니다.
```py
# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]
```
```py
# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
```
[`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py)의 계획은 두 가지 형태를 모두 사용합니다. double-stream SwiGLU 페어는 `blocks`를 직접 전달합니다. fused single-stream QKV 및 MLP 투영은 `mlp_ratio`와 `mlp_mult_factor`에 의존하기 때문에, 그 크기를 `Flux2ParallelSelfAttention.__init__`의 `Linear`에 저장합니다.
### 제외할 항목
플랜에 포함되지 않은 모듈은 모든 랭크에 복제됩니다. 정규화 레이어, 패치 및 텍스트 임베딩, 그리고 최종 `norm_out`과 `proj_out`은 제외하세요. 크기가 작아 샤딩해도 메모리 절감 효과가 크지 않고 통신 오버헤드만 증가합니다. AdaLN 모듈레이션(`img_mod`과 `txt_mod`)도 보통 크기가 작아 제외해도 무방합니다. MiniMax-H3의 `adaln_proj.linear`처럼 크기가 큰 경우 `"rowwise_split_input"`로 샤딩하세요.
### 제약 조건 및 검증
- `tp_degree`는 `config.num_attention_heads`를 나누어야 합니다.
- 각 패킹된 블록은 그 합이 아니라 *개별로* `tp_degree`로 나누어떨어져야 합니다.
[from_pretrained()](/docs/diffusers/main/ko/api/models/overview#diffusers.ModelMixin.from_pretrained)와 `enable_parallelism()`은 가중치를 로드하거나 샤딩하기 전에 둘 다 확인합니다.
눈으로 확인하는 대신 수치적으로 새로운 플랜을 검증하세요: 단일 디바이스에서 고정된 시드로 생성한 후, 텐서 병렬 환경에서 다시 생성하고 출력을 비교합니다. `"colwise"`/`"rowwise"`가 잘못 배치되어도 일반적으로 실행은 되며 그럴듯하지만 잘못된 이미지를 생성합니다.
> [!TIP]
> 비슷한 아키텍처의 기존 플랜에서 시작하세요.
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py)에는 융합된 투영이 없으므로 모든 항목은 단순한 `"colwise"` 또는 `"rowwise"`입니다. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py)에는 하나의 패킹된 row-wise 투영이 추가됩니다. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py)에는 두 가지 패킹 스타일이 모두 사용됩니다.

Xet Storage Details

Size:
11.8 kB
·
Xet hash:
c4e76533caf2eebe9451306c7d02816fa0b5cb2443519dced3b147a03a169fc7

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.