# Paralelisme tensor

[Tensor parallelism](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=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](./context_parallelism), strategi ini mengurangi memori *weight* per perangkat, yang membantu ketika model tidak muat di satu perangkat. Lihat [Distributed inference](../training/distributed_inference) untuk membandingkannya dengan strategi lainnya.

## Muat model yang terbagi

Lewatkan [TensorParallelConfig](/docs/diffusers/main/id/api/parallel#diffusers.TensorParallelConfig) ke argumen `parallel_config` dari [from_pretrained()](/docs/diffusers/main/id/api/models/overview#diffusers.ModelMixin.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](/docs/diffusers/main/id/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

Untuk menambahkan rencana ke model lain, lihat [Menambahkan paralelisme tensor ke model](#add-tensor-parallelism-to-a-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`.

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

Simpan skrip sebagai `tensor_parallel_flux2.py` dan jalankan dengan `torchrun`.

```shell
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](https://huggingface.co/black-forest-labs/FLUX.2-dev) (hanya transformer, 32B parameter, bf16) pada 4x A10G (23GB).

| tp_degree | method | waktu pemuatan | CPU puncak/rank |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64,1GB |

## Keterbatasan

Parallelisme tensor tidak berfungsi dengan beberapa opsi dan fitur pemuatan, dan Anda tidak dapat menyimpan model yang diparalelkan.

- 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()](/docs/diffusers/main/id/api/models/overview#diffusers.ModelMixin.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`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html), sehingga satu entri mencakup setiap blok. Kunci tanpa `*` berlaku untuk model itu sendiri. Path bersifat relatif terhadap model.

### Colwise dan rowwise

| Gaya | Shards | Setiap rank | Digunakan untuk |
|---|---|---|---|
| `"colwise"` | fitur output (dim 0 `weight`) | menghitung potongan dari output | `to_q`, `to_k`, `to_v`, in-projection FFN |
| `"rowwise"` | fitur input (dim 1 `weight`) | menghitung jumlah parsial | `to_out.0`, proyeksi keluaran FFN |
| `"rowwise_split_input"` | fitur input (dim 1 `weight`) | menghitung jumlah parsial dari irisan input yang direplikasi | proyeksi mandiri yang inputnya direplikasi, seperti proyeksi modulasi AdaLN yang membaca embedding timestep |

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.

```py
# 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`.

```py
"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`.

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

Rencana di [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/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 untuk juga 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()](/docs/diffusers/main/id/api/models/overview#diffusers.ModelMixin.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.

> [!TIP]
> Mulai dari rencana yang sudah ada untuk arsitektur serupa.
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py) tidak memiliki fused projections, sehingga setiap entri adalah `"colwise"` atau `"rowwise"` biasa. - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py) menambahkan satu packed row-wise projection. - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) menggunakan kedua gaya packed.