Paralelisme tensor

Tensor parallelism membagi matriks bobot model ke beberapa perangkat. Setiap perangkat menyimpan potongan kolom ("colwise") atau baris ("rowwise") dari setiap lapisan dan menghitung hasil parsial. AllReduce atau AllGather pada batas lapisan menggabungkan hasil parsial menjadi output penuh. Berbeda dengan context parallelism, strategi ini mengurangi memori weight per perangkat, yang membantu ketika model tidak muat di satu perangkat. Lihat Distributed inference untuk membandingkannya dengan strategi lainnya.

Muat model yang terbagi

Lewatkan TensorParallelConfig ke argumen parallel_config dari from_pretrained() model. tp_degree adalah jumlah perangkat untuk di-shard dan harus membagi jumlah kepala perhatian model. Kelas model harus mendefinisikan _tp_plan, yang memetakan nama modul ke gaya sharding. Untuk memeriksa apakah sebuah model mendukung tensor parallelism, cari atribut ini pada kelasnya, seperti Flux2Transformer2DModel.

Untuk menambahkan rencana ke model lain, lihat Menambahkan paralelisme tensor ke model.

Memuat dengan cara ini memecah checkpoint saat dibaca. Setiap rank hanya membaca potongan miliknya sendiri dari setiap bobot dan langsung menaruhnya di GPU-nya, sehingga tidak ada rank yang memegang bobot berukuran penuh. Memori per rank menurun seiring dengan meningkatnya 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()

Simpan skrip sebagai tensor_parallel_flux2.py dan jalankan dengan torchrun.

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

Skrip mengatur tp_degree menjadi world_size, sehingga --nproc-per-node 4 membagi transformer ke dalam 4 GPU.

Hanya transformer yang di-shard. Setiap rank memegang salinan penuh dari text encoder, sehingga setiap GPU membutuhkan memori yang cukup untuk text encoder plus bagiannya dari transformer.

Kinerja pemuatan

Memuat dengan parallel_config lebih cepat dan menggunakan lebih sedikit memori CPU per rank daripada memuat model penuh lalu memanggil enable_parallelism(). Tabel di bawah ini menunjukkan hasil untuk black-forest-labs/FLUX.2-dev (hanya transformer, 32B parameter, bf16) pada 4x A10G (23GB).

Keterbatasan

Parallelisme tensor tidak berfungsi dengan beberapa opsi dan fitur pemuatan, dan Anda tidak dapat menyimpan model yang dipecah (sharded).

• Memuat dengan parallel_config tensor-parallel tidak berfungsi dengan device_map, quantization_config, low_cpu_mem_usage=False, use_flashpack=True, atau bobot non-safetensors. Masing-masing akan memicu error alih-alih memuat checkpoint penuh secara diam-diam.
• Menggabungkan tensor parallelism dengan quantization, offloading, atau adapter LoRA juga akan memicu error, terlepas dari bagaimana model di-shard.
• Untuk memecah model yang sudah berada di memori, panggil enable_parallelism() dengan konfigurasi yang sama. Fungsi ini memuat model lengkap terlebih dahulu lalu memecahnya, sehingga setiap rank membutuhkan memori untuk checkpoint lengkap. Memuat dengan parallel_config menghindari biaya tersebut.
• Menyimpan model tensor-parallel akan memicu error di save_pretrained(). Simpan model sebelum melakukan sharding.

Menambahkan paralelisme tensor ke model

Bagian ini ditujukan bagi kontributor yang ingin menambahkan dukungan parallelisme tensor ke model Diffusers.

Tensor parallelism hanya berfungsi pada model yang mendefinisikan _tp_plan, yaitu atribut kelas yang memetakan nama modul ke gaya sharding. Menulisnya sebagian besar berarti memasangkan setiap proyeksi yang memperluas dimensi tersembunyi dengan proyeksi yang mengontraknya kembali.

Setiap kunci dapat memuat paling banyak satu *, dan bagian sebelum * harus merujuk ke nn.ModuleList, sehingga satu entri mencakup setiap blok. Kunci tanpa * berlaku untuk model itu sendiri. Path bersifat relatif terhadap model.

Colwise dan rowwise

Selalu gunakan "colwise" dan "rowwise" dalam urutan tersebut. Proyeksi "colwise" meninggalkan outputnya dalam bentuk sharded, proyeksi "rowwise" berikutnya mengonsumsi shard tersebut secara langsung, dan satu AllReduce di batas blok merekonstruksi hasilnya. Meng-shard pasangan dengan cara lain memaksa adanya gather di tengah dan menghasilkan komunikasi yang jauh lebih banyak.

Dalam attention, proyeksi Q, K, dan V per kolom memberikan setiap rank subset dari head, sehingga tp_degree harus membagi jumlah head. Duplikat encoder-stream (add_q_proj, to_add_out, ff_context) mengikuti pola yang sama dengan padanan mereka di image-stream.

Proyeksi yang digabungkan

Ketika satu Linear memuat beberapa tensor logis di sepanjang dimensi yang di-shard, potongan "colwise"/"rowwise" biasa memotong langsung melintasi konkatenasi dan menyebabkan ketidakselarasan antar bagian. Gunakan PackedColwiseParallel/PackedRowwiseParallel sebagai gantinya, yang melakukan sharding pada setiap blok yang ter-pak secara independen.

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

Lewatkan blocks, sebuah daftar bilangan bulat proporsional yang jumlahnya membagi dimensi yang dikemas. Sebagai contoh, gunakan [1, 1] untuk gate SwiGLU dan proyeksi up dengan dua bagian yang sama, atau [1, 1, 1, 3, 3] untuk proyeksi Q, K, V, gate, dan up yang digabungkan dengan mlp_ratio=3.

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

Ketika ukuran blok hanya diketahui dari config, abaikan blocks dan simpan ukuran absolut pada Linear selama __init__ sebagai _tp_packed_col_blocks atau _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(),

Rencana di transformer_flux2.py menggunakan kedua bentuk tersebut. Pasangan SwiGLU double-stream meneruskan blocks secara langsung. Proyeksi QKV dan MLP single-stream yang terfusi menyimpan ukurannya pada Linear di Flux2ParallelSelfAttention.__init__, karena ukuran tersebut bergantung pada mlp_ratio dan mlp_mult_factor.

Yang harus dikecualikan

Modul yang tidak dimasukkan ke dalam rencana akan direplikasi di setiap rank. Sisakan lapisan normalisasi, patch dan text embeddings, serta norm_out dan proj_out terakhir. Ukurannya kecil, sehingga memecahnya (sharding) hanya menghemat sedikit memori dan justru menambah komunikasi. Modulasi AdaLN (img_mod dan txt_mod) biasanya cukup kecil sehingga juga bisa dibiarkan tidak di-shard. Jika ukurannya besar, seperti adaln_proj.linear pada MiniMax-H3, shard-kan dengan "rowwise_split_input".

Kendala dan verifikasi

• tp_degree harus membagi config.num_attention_heads.
• Setiap blok yang dikemas harus secara individual dapat dibagi habis oleh tp_degree, bukan hanya jumlah totalnya.

from_pretrained() dan enable_parallelism() memeriksa keduanya sebelum memuat atau membagi (sharding) bobot apa pun.

Validasi rencana baru secara numerik, bukan dengan mata: generate dengan seed tetap pada satu perangkat, lalu ulangi di bawah tensor parallelism, dan bandingkan output-nya. Penempatan "colwise"/"rowwise" yang salah biasanya tetap berjalan dan menghasilkan gambar yang tampak masuk akal tetapi salah.

Mulai dari rencana yang sudah ada untuk arsitektur serupa.

• transformer_qwenimage.py tidak memiliki fused projections, sehingga setiap entri adalah "colwise" atau "rowwise" biasa.
• transformer_flux.py menambahkan satu proyeksi baris-berbaris yang dikemas.
• transformer_flux2.py menggunakan kedua gaya yang dikemas.