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는 하나의 packed row-wise projection을 추가합니다.

transformer_flux2.py는 두 가지 패킹 스타일을 모두 사용합니다.

Update on GitHub