Tensor paralelliği

Tensor parallelism, bir modelin ağırlık matrislerini cihazlar arasında parçalar. Her cihaz, her katman için sütun bazlı ("colwise") veya satır bazlı ("rowwise") bir dilim tutar ve kısmi bir sonuç hesaplar. Katman sınırındaki bir AllReduce veya AllGather, kısmi sonuçları birleştirerek tam çıktıyı oluşturur. context parallelism'den farklı olarak, cihaz başına weight belleğini azaltır; bu da bir modelin tek bir cihaza sığmadığı durumlarda işe yarar. Diğer stratejilerle karşılaştırmak için Distributed inference bölümüne bakın.

Parçalı bir model yükleyin

Modelin from_pretrained() metodunun parallel_config argümanına bir TensorParallelConfig geçirin. tp_degree, parçalanacak cihaz sayısıdır ve modelin dikkat başlıklarının sayısına tam bölünmelidir. Modelin sınıfı, modül adlarını bir parçalama stiline eşleyen bir _tp_plan tanımlamalıdır. Bir modelin tensör paralelliğini destekleyip desteklemediğini kontrol etmek için, Flux2Transformer2DModel gibi, sınıfında bu özniteliği arayın.

Başka bir modele plan eklemek için Bir modele tensör paralelliği ekleme bölümüne bakın.

Bu yöntemle yüklenirken, kontrol noktası okunurken parçalara ayrılır. Her rank, her ağırlığın yalnızca kendi dilimini okur ve doğrudan kendi GPU'suna yerleştirir; böylece hiçbir rank tam boyutlu bir ağırlık tutmaz. Rank başına bellek, tp_degree arttıkça azalır.

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()

Betik dosyasını tensor_parallel_flux2.py olarak kaydedin ve torchrun ile başlatın.

torchrun --nproc-per-node 4 tensor_parallel_flux2.py

Betik, tp_degree değerini world_size olarak ayarladığından, --nproc-per-node 4 komutu dönüştürücüyü 4 GPU arasında paylaştırır.

Sadece transformer parçalanır. Her rank, metin kodlayıcısının tam bir kopyasını tutar; bu nedenle her GPU'nun, metin kodlayıcı ile transformer'ın kendine düşen payı için yeterli belleğe sahip olması gerekir.

Yükleme performansı

parallel_config ile yükleme, tam modeli yükleyip ardından enable_parallelism() çağrılmasından daha hızlıdır ve rank başına daha az CPU belleği kullanır. Aşağıdaki tablo, 4x A10G (23GB) üzerinde black-forest-labs/FLUX.2-dev (yalnızca transformer, 32B parametre, bf16) için sonuçları göstermektedir.

Kısıtlamalar

Tensor parallelism, bazı yükleme seçenekleri ve özelliklerle çalışmaz ve bölünmüş bir modeli kaydedemezsiniz.

• parallel_config ile tensor-parallel yükleme, device_map, quantization_config, low_cpu_mem_usage=False, use_flashpack=True veya safetensors olmayan ağırlıklarla çalışmaz. Bunların her biri, tam checkpoint'i sessizce yüklemek yerine bir hata fırlatır.
• Gergilik paralelliğinin, modelin nasıl bölündüğüne bakılmaksızın, kuantizasyon, offloading veya LoRA adaptörleriyle birleştirilmesi de bir hata oluşturur.
• Bellekte zaten bulunan bir modeli parçalamak için aynı yapılandırmayla enable_parallelism() çağrılır. Bu yöntem önce tam modeli yükler ve ardından parçalar; bu nedenle her rank, tam checkpoint için bellek gerektirir. Bu maliyeti önlemek için yükleme sırasında parallel_config kullanılır.
• Tensor-parallel bir modelin kaydedilmesi save_pretrained() içinde bir hata oluşturur. Modeli parçalara ayırmadan önce kaydedin.

Bir modele tensör paralelliği ekleyin

Bu bölüm, bir Diffusers modeline tensör paralelliği desteği eklemek isteyen katkıda bulunanlar içindir.

Tensor paralelliği, yalnızca modül adlarını bir parçalama stiline eşleyen _tp_plan sınıf özniteliğini tanımlayan modellerde çalışır. Bu özniteliği yazmak, çoğunlukla gizli boyutu genişleten her projeksiyonu, onu geri daraltan projeksiyonla eşleştirmek anlamına gelir.

Her anahtar en fazla bir * içerebilir ve * işaretinden önceki bölüm bir nn.ModuleList'ye çözümlenmelidir, böylece bir girdi tüm blokları kapsar. * içermeyen bir anahtar, modelin kendisine uygulanır. Yollar modele göre göreli olarak tanımlanır.

Colwise ve rowwise

"colwise" ve "rowwise" değerlerini her zaman bu sırayla eşleştirin. "colwise" projeksiyonu çıktısını parçalı (sharded) bırakır, ardından gelen "rowwise" projeksiyonu bu parçayı doğrudan tüketir ve blok sınırındaki tek bir AllReduce işlemi sonucu yeniden oluşturur. Bu çiftin başka bir şekilde parçalanması, ortada bir toplama (gather) işlemi zorunlu kılar ve çok daha fazla iletişim gerektirir.

Attention'da sütun bazlı Q, K ve V projeksiyonları her rank'a başların bir alt kümesini verir, bu nedenle tp_degree baş sayısına tam bölünmelidir. Encoder-stream kopyaları (add_q_proj, to_add_out, ff_context), image-stream karşılıklarıyla aynı deseni izler.

Birleştirilmiş projeksiyonlar

Tek bir Linear, parçalanan boyut boyunca birden fazla mantıksal tensörü paketlediğinde, düz "colwise"/"rowwise" dilimleri birleştirmenin üzerinden doğrudan geçer ve parçaları hizasız bırakır. Bunun yerine, her paketlenmiş bloğu bağımsız olarak parçalayan PackedColwiseParallel/PackedRowwiseParallel kullanın.

# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel

blocks, paketlenmiş boyutu bölen oransal tam sayılardan oluşan bir liste, geçirin. Örneğin, eşit yarımlara sahip bir SwiGLU gate ve up projection için [1, 1] veya mlp_ratio=3 ile birleşik Q, K, V, gate ve up projection için [1, 1, 1, 3, 3] kullanın.

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

Blok boyutları yalnızca yapılandırmadan biliniyorsa, blocks parametresini atlayın ve mutlak boyutları __init__ sırasında Linear üzerinde _tp_packed_col_blocks veya _tp_packed_row_blocks olarak saklayın.

# 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 içindeki plan her iki biçimi de kullanır. Çift akışlı SwiGLU çifti blocks değerini doğrudan iletir. Birleşik tek akışlı QKV ve MLP projeksiyonu, boyutları mlp_ratio ve mlp_mult_factor değerlerine bağlı olduğu için, bu boyutları Flux2ParallelSelfAttention.__init__ içindeki Linear nesnesinde saklar.

Dışarıda bırakılacaklar

Plan dışı bırakılan modüller her sırada kopyalanır. Normalizasyon katmanlarını, yama ve metin gömme katmanlarını ile son norm_out ve proj_out katmanlarını dışarıda bırakın. Bunlar küçük olduğu için bunları bölümlere ayırmak az miktarda bellek tasarrufu sağlar ve iletişim yükü ekler. AdaLN modülasyonu (img_mod ve txt_mod) genellikle dışarıda bırakılacak kadar küçüktür. MiniMax-H3’ün adaln_proj.linear gibi büyük olduğunda, "rowwise_split_input" ile bölümlere ayırın.

Kısıtlamalar ve doğrulama

• tp_degree, config.num_attention_heads'i tam olarak bölmelidir.
• Her paketlenmiş blok, yalnızca toplamları değil, tek tek tp_degree değerine bölünebilir olmalıdır.

from_pretrained() ve enable_parallelism(), herhangi bir ağırlık yüklenmeden veya bölünmeden önce her iki kontrolü de yapar.

Yeni bir planı gözle değil, sayısal olarak doğrulayın: sabit bir tohumla tek bir cihazda üretim yapın, ardından tensor paralelliği altında tekrar üretin ve çıktıları karşılaştırın. Yanlış yerleştirilmiş bir "colwise"/"rowwise" genellikle yine de çalışır ve makul görünen ancak yanlış bir görüntü üretir.

Benzer bir mimari için mevcut bir plandan başlayın.

• transformer_qwenimage.py dosyasında füzyonlu projeksiyonlar bulunmadığından, her giriş düz bir "colwise" veya "rowwise" değeridir.
• transformer_flux.py tek bir paketlenmiş satır bazlı projeksiyon ekler.
• transformer_flux2.py her iki paketli stili de kullanır.