diffusers/main/ko/optimization/tensor_parallelism.md

Tensor 병렬성

Tensor parallelism은 모델의 가중치 행렬을 여러 디바이스에 걸쳐 분할합니다. 각 디바이스는 각 레이어의 열 방향("colwise") 또는 행 방향("rowwise") 슬라이스를 보유하며 부분 결과를 계산합니다. 레이어 경계에서의 AllReduce 또는 AllGather 연산은 부분 결과를 결합하여 전체 출력을 생성합니다. context parallelism과 달리, 이는 디바이스당 weight 메모리를 줄여주므로 모델이 단일 디바이스에 맞지 않을 때 유용합니다. 다른 전략과의 비교를 위해 Distributed inference를 참조하세요.

Load a sharded model

Pass a TensorParallelConfig to the parallel_config argument of the model's 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.

다른 모델에 플랜을 추가하려면 모델에 텐서 병렬성 추가를 참조하세요.

이 방식으로 로드하면 체크포인트가 읽히는 동안 샤드됩니다. 각 랭크는 각 가중치의 자기 자신의 슬라이스만 읽고 이를 해당 GPU에 직접 배치하므로, 어떤 랭크도 전체 크기의 가중치를 보유하지 않습니다. 랭크당 메모리는 tp_degree가 증가함에 따라 감소합니다.

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으로 실행하세요.

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 (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()에서 오류를 발생시킵니다. 샤딩하기 전에 모델을 저장하세요.

모델에 텐서 병렬성 추가

이 섹션은 Diffusers 모델에 텐서 병렬성(tensor parallelism) 지원을 추가하려는 기여자를 위한 것입니다.

Tensor parallelism은 _tp_plan이라는 클래스 속성을 정의한 모델에서만 작동합니다. 이 속성은 모듈 이름을 샤딩 스타일에 매핑합니다. 이 속성을 작성하는 것은 대부분 확장하는 투영과 축소하는 투영을 짝짓는 것을 의미합니다.

각 키에는 *가 최대 하나만 포함될 수 있으며, * 앞부분은 nn.ModuleList로 해석되어야 하므로 하나의 항목이 모든 블록을 커버합니다. *가 없는 키는 모델 자체에 적용됩니다. 경로는 모델 기준 상대 경로입니다.

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을 사용하세요.

# 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]을 사용합니다.

"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),

blocks가 config에서만 알려진 경우, blocks를 생략하고 대신 __init__ 중에 Linear에 절대 크기를 _tp_packed_col_blocks 또는 _tp_packed_row_blocks로 저장합니다.

# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]

# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),

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()와 enable_parallelism()은 가중치를 로드하거나 샤딩하기 전에 둘 다 확인합니다.

눈으로 확인하는 대신 수치적으로 새로운 플랜을 검증하세요: 단일 디바이스에서 고정된 시드로 생성한 후, 텐서 병렬 환경에서 다시 생성하고 출력을 비교합니다. "colwise"/"rowwise"가 잘못 배치되어도 일반적으로 실행은 되며 그럴듯하지만 잘못된 이미지를 생성합니다.

비슷한 아키텍처의 기존 플랜에서 시작하세요.

• transformer_qwenimage.py에는 융합된 투영이 없으므로 모든 항목은 단순한 "colwise" 또는 "rowwise"입니다. - transformer_flux.py에는 하나의 패킹된 row-wise 투영이 추가됩니다. - transformer_flux2.py에는 두 가지 패킹 스타일이 모두 사용됩니다.