Paralelisme tensor untuk inferensi

Tensor parallelism memotong lapisan model menjadi beberapa bagian agar beberapa akselerator perangkat keras dapat memprosesnya secara bersamaan. Hal ini memungkinkan Anda menjalankan model yang melebihi kapasitas memori satu GPU dan mencapai throughput yang lebih tinggi. Anda memerlukan komunikasi intra-node yang cepat karena GPU bertukar hasil parsial di setiap lapisan.

Model mendukung tensor parallelism jika config-nya mendefinisikan base_model_tp_plan. Periksa model yang dimuat dengan properti supports_tp_plan.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

Panduan ini membahas cara mengaktifkan tensor parallelism di Transformers serta strategi partisi yang tersedia.

Partisi model

Konfigurasikan jumlah perangkat paralel tensor dengan tp_size di DistributedConfig.

Set DistributedConfig(tp_size=N) untuk menggunakan rencana yang telah didefinisikan sebelumnya oleh model.
Definisikan tp_plan secara manual dan teruskan ke DistributedConfig bersama dengan tp_size.

Anda juga dapat mengatur tp_plan="auto" untuk meminta rencana yang sudah didefinisikan secara eksplisit. Ketika tp_size diabaikan dan tp_plan telah diatur, tp_size diturunkan dari WORLD_SIZE dibagi dengan ukuran paralel lainnya. Menyalurkan tp_plan langsung ke from_pretrained() sudah tidak didukung (deprecated) dan akan dihapus pada v5.18.

auto plan

rencana manual

Strategi partisi

ALL_PARALLEL_STYLES, sebuah ParallelInterface yang didefinisikan di distributed/parallel_utils.py, memetakan setiap nama strategi yang dapat Anda gunakan dalam sebuah rencana ke instance strategi yang telah dikonfigurasi. Anda tidak berinteraksi dengannya secara langsung untuk melakukan sharding model, tetapi ini adalah daftar resmi nama-nama yang tersedia. distributed/tensor_parallel.py mendaftarkan strategi-strategi tensor parallel ke dalamnya:

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

Setiap strategi tensor parallel adalah subclass dari TensorParallelLayer di file tersebut. Strategi expert parallel, seperti ep_dispatch_experts, berada di distributed/expert_parallel.py, subclass dari ExpertParallelLayer yang memiliki hook yang sama, dan terdaftar ke dalam ALL_PARALLEL_STYLES yang sama. Dua strategi yang paling sering Anda gunakan, ColwiseParallel dan RowwiseParallel, menerima penempatan input_layouts dan output_layouts, yang memungkinkan satu kelas mencakup beberapa nama plan. colwise meninggalkan output-nya ter-shard pada dimensi terakhir untuk dikonsumsi oleh lapisan rowwise berikutnya, sedangkan colwise_gather_output melakukan all-gather untuk mengembalikannya menjadi tensor penuh.

Tabel di bawah ini menjelaskan setiap strategi.

Strategi	Deskripsi
`ColwiseParallel`	Membagi (shard) bobot linear pada dimensi fitur output (`Shard(0)` untuk bobot `nn.Linear` 2D) dan membagi bias 1D. Mendistribusikan ulang input ke `Replicate()`, kemudian menempatkan output di `output_layouts`.
`RowwiseParallel`	Shards bobot linear pada dimensi fitur inputnya (`Shard(-1)`) dan mereplikasi bias, yang ditambahkan setelah reduksi. Forward pada setiap rank menghasilkan output `Partial()` yang menyimpan bagiannya dari jumlah tersebut, dan `transform_output_post_forward` mendistribusikan ulang bagian itu ke `output_layouts`. Mereduksi ke `Replicate()` memicu all-reduce, dan mereduksi ke `Shard(1)` memicu reduce-scatter.
`EmbeddingParallel`	Membagi `nn.Embedding` dengan token id yang direplikasi. `shard_dim=0` (`embedding_rowwise`, pilihan yang umum) membagi kosakata dan melakukan all-reduce pada hasil pencarian parsial dari setiap rank menjadi output lengkap. `shard_dim=1` (`embedding_colwise`) membagi dimensi tersembunyi dan melakukan all-gather pada potongan output dari setiap rank.
`PackedColwiseParallel`	Varian dari `ColwiseParallel` untuk bobot yang digabungkan (fused), misalnya `up_proj` dan `gate_proj` yang dikemas menjadi `gate_up_proj`. Gunakan `split_factor` ketika sebuah bobot mengemas jumlah lapisan selain dua.
`PackedRowwiseParallel`	Padanan per-baris, untuk bobot yang dikemas di sepanjang dimensi akhir. Mereplikasi parameter 1D.
`SequenceParallel`	Menggandakan parameter modul dan memecah inputnya pada `sequence_dim` (default `1`). Digunakan untuk norm yang beroperasi per-token, seperti `LayerNorm` dan `RMSNorm`.
`ReplicatedWithGradAllReduce`	Menggandakan sebuah parameter tetapi melakukan all-reduce pada gradiennya. Diperlukan untuk norm yang berada di antara lapisan column-wise dan row-wise serta melakukan normalisasi sepanjang sumbu yang di-shard, di mana setiap rank hanya melihat head miliknya sendiri.
`AllReduceParallel`	Melakukan all-reduce pada output forward `Partial()` dari sebuah modul menjadi `Replicate()`. Gunakan ini sebagai titik sinkronisasi untuk modul yang komputasinya berakhir dengan penjumlahan parsial.
`MlaKvAProjParallel`	Memecah output `kv_a_proj_with_mqa` dari perhatian MLA bergaya DeepSeek-V2 dan melakukan all-reduce pada gradien bagian RoPE, yang melewati `kv_b_proj` dan jika tidak akan mempertahankan gradien parsial. Memerlukan `qk_rope_head_dim` dalam konfigurasi model.
`MoeExpertsParallel`	Ekspert MoE paralel tensor. Melakukan all-reduce pada output ekspert di arah maju dan menambahkan all-reduce di arah mundur untuk hidden states dan bobot routing.
`MoeIdentityParallel`	Membagi input dari ahli nol atau identitas dengan ukuran mesh, sehingga membatalkan all-reduce yang diterapkan oleh `moe_tp_experts` di hilir.
`LocalParamsParallel`	Menjalankan forward sebuah modul dengan parameter terbagi miliknya sendiri sebagai tensor lokal. Gunakan ini pada pemilik `nn.Parameter` polos, seperti attention sinks, yang dipecah oleh entri rencana miliknya sendiri.
`MoEParamShard`	Shards bobot pakar MoE pada penempatan tertentu. Mendukung nama `grouped_gemm`, di mana `shards_expert_dim=True` juga menulis ulang `module.num_experts` menjadi jumlah pakar per-rank.
`EpRouterParallel`	Masks skor router untuk pakar non-lokal dan memetakan ulang ID pakar global menjadi ID lokal, sehingga setiap rank hanya menjalankan pakar yang dimilikinya. Memerlukan `num_experts` yang dapat dibagi habis oleh ukuran mesh.
`RouterParallelMegaMoe`	Router variant untuk DeepGEMM Mega MoE, yang mendispatch para ahli di dalam kernel yang menginginkan output router tetap utuh.
`MoeTensorParalellMegaMoeExperts`	Lapisan ahli khusus inferensi untuk DeepGEMM Mega MoE. Melewati sinkronisasi gradien dan meneruskan grup proses ke modul agar kernel dapat menyiapkan buffer bersama pada forward pertama.

Strategi yang dikemas

Weight packing menggabungkan beberapa lapisan linear menjadi satu lapisan yang lebih besar. Strategi PackedColwiseParallel dan PackedRowwiseParallel melakukan sharding pada bobot yang sudah di-packing dengan benar. Strategi ColwiseParallel atau RowwiseParallel dasar melakukan sharding pada bobot yang sudah di-packing dengan tidak benar.

Contoh di bawah ini mengemas up_proj dan gate_proj ke dalam satu modul gate_up_proj dan memerlukan strategi packed_rowwise untuk mempartisi gate_up_proj.

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

Gunakan perkalian matriks batch pada jalur forward untuk menghitung keluaran modul gate_up_proj.

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

Shard biasa membagi dimensi tersebut menjadi blok-blok yang berurutan, sehingga sebuah rank akan menerima ekor dari gate dan kepala dari up alih-alih potongan dari masing-masing. Strategi yang dikemas menggunakan _StridedShard dengan split_factor (2 secara default) untuk menginterleaving pembagian, sehingga setiap rank mendapatkan potongan yang sesuai dari kedua bagian agar chunk tetap sejajar setelah sharding.

Strategi partisi kustom

Warisi dari TensorParallelLayer di distributed/tensor_parallel.py untuk membuat strategi partisi kustom. Override hanya hook yang dibutuhkan oleh strategi Anda, karena setiap hook memiliki default no-op.

Hook	Tujuan
`validate_param`	Tolak parameter yang tidak dapat dipecah oleh strategi ini, sebelum bobot dimuat.
`shard_param`	Ganti satu parameter dengan placeholder `DTensor` agar loader mengetahui shard mana yang milik rank ini.
`transform_inputs_pre_forward`	Distribusikan ulang input modul ke tata letak yang diharapkan oleh forward-nya.
`context_around_forward`	Bungkus forward dalam context manager, misalnya untuk mengekspos tensor lokal ke kernel.
`transform_output_post_forward`	Redistribusikan atau kurangi output modul.
`should_use_local_tensors`	Laporkan bahwa forward modul ini memerlukan tensor polos, bukan `DTensor`.
`install_forward`	Ganti `module.forward` secara langsung. Override ini hanya jika hook di atas tidak cukup, seperti yang dilakukan `ReplicatedWithGradAllReduce` untuk mendaftarkan hook backward.

Contoh di bawah ini membahas versi ringkas dari ColwiseParallel.

Turunkan dari TensorParallelLayer dan simpan penempatan yang digunakan oleh strategi tersebut. Kelas dasar tidak mendefinisikan __init__, sehingga tidak ada yang perlu dipanggil dengan super().

class ColwiseParallel(TensorParallelLayer):
    def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
        self.input_layouts = input_layouts or Replicate()
        self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
        self.use_local_output = use_local_output

Implementasikan shard_param untuk membungkus satu parameter sebagai placeholder DTensor. Karena berjalan pada meta tensor, distribute_tensor hanya membangun metadata dan tidak memindahkan data. Sampaikan src_data_rank=None karena belum ada tensor penuh untuk disebarkan.

def shard_param(self, module, param, mesh):
    meta = module._parameters.get(param)
    if meta is None:
        return
    # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
    placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
    module._parameters[param] = torch.nn.Parameter(
        distribute_tensor(meta, mesh, [placement], src_data_rank=None),
        requires_grad=meta.requires_grad,
    )

Implementasikan transformasi input dan output. install_forward memanggil transformasi tersebut di sekitar forward asli modul, sehingga mereka hanya perlu memindahkan tensor antar layout. Partisi searah kolom mengharapkan input yang direplikasi dan menghasilkan output yang di-shard pada dimensi terakhir.

def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
    x = args[0]
    if not isinstance(x, DTensor):
        x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
    if x.placements != (Replicate(),):
        x = x.redistribute(placements=[Replicate()])
    return (x,) + args[1:], kwargs

def transform_output_post_forward(self, module, output, mesh):
    if not isinstance(output, DTensor):
        output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
    if output.placements != (self.output_layouts,):
        output = output.redistribute(placements=[self.output_layouts])
    return output.to_local() if self.use_local_output else output

ColwiseParallel yang disertakan menambahkan jalur cepat di atas ini yang melewati bolak-balik DTensor untuk inferensi nn.Linear biasa dan untuk modul terkuantisasi yang membutuhkan tensor lokal. Baca sumbernya sebelum menyalinnya jika strategi Anda membutuhkan perlakuan yang sama.

Daftarkan strategi agar tp_plan dapat menamainya. Pendaftaran menerima instance strategi, sehingga penempatan yang Anda lewati ke __init__ akan disimpan bersama namanya. Perhatikan () di bawah ini.

import torch

from transformers import AutoModelForCausalLM, DistributedConfig
from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES

ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise_custom",
    ...
}
distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

Benchmarks

Parallelisme tensor secara signifikan mempercepat inferensi, terutama untuk ukuran batch besar atau sekuens panjang.

Grafik ini menunjukkan percepatan yang diharapkan untuk satu kali forward pass pada Llama dengan panjang sekuens 512.

Implementasi desain

Transformers mengimplementasikan parallelism tensor dengan cara yang tidak bergantung pada framework. Ia mengandalkan DeviceMesh dan DTensor dari torch.distributed untuk menyediakan antarmuka yang sederhana dan dapat diperluas.

DeviceMesh

DeviceMesh membuat grid perangkat multi-dimensi yang saling berkomunikasi. Strategi paralelisasi yang berbeda memerlukan pola komunikasi yang berbeda. Buat DeviceMesh dengan beberapa sub-mesh untuk menangani pola-pola tersebut.

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

Sebagian besar strategi paralelisasi torch.distributed diterapkan pada mesh itu sendiri atau sub-mesh-nya. Mesh secara otomatis menangani pola komunikasi.

DTensor

DTensor (Distributed Tensor) menangani logika distribusi di atas operasi tensor biasa. Bobot model di bawah tensor parallelism disimpan sebagai DTensor, yang memungkinkan strategi mendeskripsikan komunikasi sebagai perubahan layout alih-alih kolektif eksplisit.

Atribut placements memberi tahu PyTorch bagaimana sebuah tensor ditata di antara perangkat dalam DeviceMesh. Atribut ini menerima nilai-nilai berikut:

Shard(dimension) membagi DTensor di sepanjang dimensi yang ditentukan pada DeviceMesh tempat tensor tersebut dibuat. Partisi searah kolom membagi dimensi fitur keluaran dari bobot dan satu-satunya dimensi dari bias.

weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension

Partisi per baris memecah dimensi fitur input dan mereplikasi bias, karena bias ditambahkan sekali setelah reduksi, bukan di setiap rank.

weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs

Replicate() mereplikasi DTensor di seluruh DeviceMesh, membuat salinan penuh tensor di setiap perangkat.
Partial() menandai tensor sebagai menunggu reduksi. Output forward dari layer row-wise adalah Partial(), dan redistribusinya ke Replicate() adalah yang memicu all-reduce.

Sumber daya

Bagian Ultra-Scale Playbook tentang tensor parallelism menyediakan lebih banyak detail.
Periksa panduan expert parallelism jika Anda menggunakan model mixture-of-experts (MoE). Model-model ini mendukung tensor parallelism dan expert parallelism.
Baca artikel blog Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand untuk ringkasan singkat tentang tensor parallelism dan pelajari bagaimana setup column dan row parallel berbeda.
Lihat panduan pelatihan Tensor parallelism untuk mempelajari cara menggunakannya dalam pengaturan pelatihan.

Update on GitHub