Paralelisme pakar

Expert parallelism adalah strategi paralelisasi untuk model mixture-of-experts (MoE). Lapisan feedforward dari setiap expert berada di akselerator perangkat keras yang berbeda. Router meneruskan token ke expert yang sesuai dan mengumpulkan hasilnya. Pendekatan ini memungkinkan model diskalakan ke jumlah parameter yang jauh lebih besar tanpa meningkatkan biaya komputasi, karena setiap token hanya mengaktifkan beberapa expert.

DistributedConfig

Aktifkan paralelisme pakar dengan kelas DistributedConfig dan argumen ep_size. Sebagian besar model MoE secara default menggunakan token dispatch, sehingga ep_size dapat diatur secara independen dari tp_size. Beberapa model, seperti Llama 4 dan Gemma 4, masih secara default menggunakan masking dan all-reduce ("ep_router" dan "moe_tp_experts" di model.ep_plan). Masking juga tersedia pada model apa pun dengan override ep_plan, dan memerlukan ep_size=tp_size agar setiap rank dalam grup pakar menerima token yang sama.

import os

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.distributed.configuration_utils import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=int(os.environ["WORLD_SIZE"]),
    ep_size=int(os.environ["WORLD_SIZE"]),
)

model = AutoModelForCausalLM.from_pretrained(
    "openai/gpt-oss-120b",
    distributed_config=distributed_config,
)

Setiap model MoE mendefinisikan dua rencana dalam konfigurasinya: base_model_tp_plan untuk modul padat dan base_model_ep_plan untuk para ahli. Keduanya dipaparkan pada model yang dimuat sebagai model.tp_plan dan model.ep_plan. Dengan tp_size > 1 dan ep_size > 1, keduanya berlaku: rencana paralel tensor memecah perhatian dan MLP padat, sedangkan rencana paralel ahli memecah para ahli. Aturan EP memiliki prioritas lebih tinggi daripada aturan TP untuk modul yang sama, sehingga bobot ahli hanya dipecah sekali, oleh rencana EP. Dalam rencana EP, gaya GroupedGemmParallel membagi bobot ahli sepanjang dimensi ahli sehingga setiap rank hanya memuat ahli lokalnya, dan ep_router memfilter ahli yang berada di rank lain sebelum all-reduce menggabungkan keluaran para ahli.

tp_plan hanya diterapkan ketika tp_size > 1, dan ep_plan hanya ketika ep_size > 1. Dengan TP diaktifkan dan EP dinonaktifkan, rencana TP penuh diterapkan, termasuk aturan expert.

Aturan forward ahli dalam ep_plan memilih cara token mencapai para ahli:

enable_expert_parallel=True tanpa ep_size sudah tidak didukung dan mengeluarkan FutureWarning. Ini hanya memecah para ahli, seperti ep_size=tp_size, tp_plan={}.

Jalankan skrip inferensi Anda dengan torchrun. Jumlah proses harus sama dengan tp_size * fsdp_size * pp_size, dan ep_size harus membagi jumlah pakar secara merata.

torchrun --nproc-per-node 8 your_script.py

Mengganti rencana

Lewati tp_plan={...} atau ep_plan={...} ke DistributedConfig untuk menimpa aturan individual dari rencana yang telah ditentukan sebelumnya. Aturan yang tidak ditentukan akan dipertahankan, dan rencana yang digabungkan disimpan pada model. Setiap kunci harus cocok dengan modul, parameter, atau entri rencana yang ada; jika tidak, pemuatan akan memicu ValueError sebelum apa pun dipecah. Gunakan jalur lengkap sebagaimana dilihat dari model yang dimuat, sehingga model.layers.* untuk LM kausal dan layers.* untuk model dasarnya.

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    tp_plan={"model.layers.*.self_attn.q_proj": "colwise_rep"},
    ep_plan={"model.layers.*.mlp.experts.down_proj": "grouped_gemm"},
)

Menyediakan rencana tidak menyimpulkan ukuran paralel: atur tp_size dan ep_size secara eksplisit.

Sebagian besar model MoE secara default menggunakan "ep_dispatch_experts". Untuk menggunakan masking dan all-reduce sebagai gantinya, atur ep_size=tp_size dan timpa aturan forward untuk router maupun expert (nama modul router bergantung pada model, misalnya mlp.router pada gpt-oss):

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)

Sebaliknya, timpa aturan forward expert dari model yang rencananya menggunakan masking dengan "ep_dispatch_experts" untuk menggunakan token dispatch. Aturan router kemudian diabaikan, karena dispatch membutuhkan id expert global untuk menemukan pemilik masing-masing expert.

Token dispatch

Dengan token dispatch, setiap rank melatih pada bagian batch-nya sendiri. Pada setiap lapisan MoE, sebuah rank merutekan token-nya, mengirim setiap pasangan (token, expert) ke rank yang memiliki expert tersebut melalui all-to-all, menjalankan expert lokalnya pada apa yang diterimanya, menerima hasilnya kembali melalui all-to-all kedua, dan menggabungkannya dengan bobot routing. Hanya aktivasi yang dirutekan dan output expert yang berpindah, dan tidak ada rank yang menghitung expert untuk token yang tidak dimilikinya.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=1,
    fsdp_size=8,
    ep_size=4,
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

Dengan tp_size=1, ep_size harus membagi habis fsdp_size dan jumlah pakar, dan perhatian tidak dibatasi oleh num_key_value_heads. Untuk sisa model:

• Parameter di luar para ahli dipecah dengan FSDP2 di sepanjang fsdp, yang mengurangi gradien mereka.
• Para ahli di-shard di ep dan, ketika efsdp_size = fsdp_size * tp_size // ep_size lebih besar dari satu, juga di-shard FSDP di efsdp. Mereka selalu dibungkus FSDP, sehingga fsdp_mixed_precision dan fsdp_cpu_offload juga berlaku untuk mereka dan save_pretrained() mengumpulkannya seperti parameter lainnya.
• Grup paralel ahli memegang ep_size / tp_size batch, sehingga gradien dari sebuah ahli adalah jumlah dari sebanyak itu batch. Reduksi efsdp membagi dengan fsdp_size alih-alih ukuran grupnya, yang menghasilkan rata-rata per-batch yang sama dengan yang diambil FSDP2 untuk modul padat.
• Setiap pakar lokal juga memproses satu baris nol pad per lapisan. Rank yang pakarnya tidak menerima token apa pun tetap bergabung dalam reverse all-to-all dan pengurangan gradien pakar.

Bagaimana ukuran-ukuran tersebut digabungkan

• tp_size * fsdp_size adalah jumlah proses. ep_size tidak menambah apa pun: ia hanya mengelompokkan ulang rank yang sama untuk bobot pakar saja.
• ep_size memotong daftar expert menjadi ep_size blok. Setiap rank menghitung num_experts / ep_size expert, dan ep_size rank berurutan memegang satu set lengkap. Set rank tersebut adalah grup tempat all-to-all dijalankan.
• efsdp_size = fsdp_size * tp_size / ep_size adalah jumlah salinan lengkap dari set expert. Rank pada posisi yang sama di salinan yang berbeda melakukan sharding terhadap expert tersebut untuk memori dan merata-ratakan gradien mereka, seperti yang dilakukan FSDP untuk modul dense.
• Batch yang dipegang oleh sebuah rank hanya bergantung pada fsdp. Rank-rank yang berurutan membentuk grup TP dan menerima batch yang sama; grup-grup fsdp_size menerima batch yang berbeda.

Delapan proses, tp_size=2, fsdp_size=4, delapan ahli:

rank         0      1       2      3       4      5       6      7
batch       [====B0====]   [====B1====]   [====B2====]   [====B3====]     one batch per TP pair
tp           0      1       0      1       0      1       0      1

ep_size=2   E0-3   E4-7    E0-3   E4-7    E0-3   E4-7    E0-3   E4-7     group = a TP pair,   efsdp_size=4
ep_size=4   E0E1   E2E3    E4E5   E6E7    E0E1   E2E3    E4E5   E6E7     group = two pairs,   efsdp_size=2
ep_size=8   E0     E1      E2     E3      E4     E5      E6     E7       group = all ranks,   efsdp_size=1

Dua angka dapat disimpulkan dari gambar:

• Di dalam grup EP, setiap token ada tp_size kali, satu kali per rank dari pasangan yang memegang batch-nya. Hal ini tidak bergantung pada ep_size.
• Grup EP memegang ep_size / tp_size batch yang berbeda. Ini adalah jumlah yang dijumlahkan oleh gradien seorang expert.

Dengan paralelisme tensor

Setel tp_size > 1 untuk memecah modul padat dengan rencana TP, sementara para ahli menggunakan dispatch. Pada delapan proses:

distributed_config = DistributedConfig(
    tp_size=2,
    fsdp_size=4,
    ep_size=4,
)

Setiap pasangan rank TP menerima batch yang sama, karena paralelisme tensor mereplikasi aktivasi di dalam pasangan tersebut. Jika kedua rank mendispatch semua token mereka, rank yang memiliki sebuah expert akan menerima setiap token dua kali, menghitungnya dua kali, dan gradien bobotnya akan menjadi dua kali lipat. Expert utuh berada di satu rank, sehingga duplikat tidak dapat dipecah berdasarkan bobot, dan pemiliknya biasanya adalah rank lain, sehingga tidak dapat diselesaikan berdasarkan kepemilikan seperti yang dilakukan masking. Oleh karena itu, pasangan tersebut membagi baris: setiap rank TP mendispatch 1 / tp_size token yang tidak tumpang tindih, hasil kembali ke rank yang mengirimkannya, dan all-reduce di atas pasangan dari separuh yang di-nolkan memulihkan output yang direplikasi yang diharapkan oleh lapisan berikutnya. Pembagian dilakukan berdasarkan tp_size, bukan ep_size, karena hanya rank yang memegang batch yang dapat mengirimkannya. Grup expert membentang di empat rank dan setiap expert di-shard FSDP di efsdp_size = 2 rank, sementara grup FSDP trunk membentang di empat rank. Kendala TP biasa dari model, seperti keterbagian kepala perhatian, masih berlaku untuk modul padat. Irisan token mungkin tidak rata atau kosong, termasuk selama decoding token tunggal.

Token dispatch belum dapat dikombinasikan dengan pipeline parallelism; gunakan pp_size=1 (belum diuji) Konfigurasi ini masing-masing menggunakan delapan GPU:

Trainer belum memperhitungkan dispatch token: penghitungan batch dan token mengasumsikan tata letak all-reduce, di mana rank dari grup TP berbagi satu batch dan terdapat fsdp_size shard data-parallel. Dukungan Trainer untuk dispatch akan hadir pada pembaruan berikutnya.

Menggabungkan dengan FSDP2

Tensor dan expert parallelism memecah bobot di sepanjang tp, tetapi state optimizer dan modul tanpa aturan masih direplikasi di setiap rank dalam grup, yang membatasi seberapa besar model yang dapat dilatih. Tambahkan FSDP2 pada dimensi mesh kedua dengan fsdp_size. Dengan masking dan all-reduce, pertahankan ep_size=tp_size dan lewati `ep_plan={“layers..mlp.gate”: “ep_router”, “layers..mlp.experts”: “moe_tp_experts”}“.

from transformers import AutoModelForCausalLM
from transformers.distributed import DistributedConfig

distributed_config = DistributedConfig(
    tp_size=4,
    ep_size=4,  # expert parallel size, must match tp_size with masking and all-reduce
    fsdp_size=2,  # data parallel shards
    ep_plan={
        "model.layers.*.mlp.gate": "ep_router",
        "model.layers.*.mlp.experts": "moe_tp_experts",
    },
)
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-30B-A3B", distributed_config=distributed_config)

Muat model seperti biasa, lalu latih dengan Trainer. Proses ini mengambil norma gradien di kedua mesh dan memberikan grup parameter optimizer yang terpisah untuk masing-masing mesh. save_model() mengumpulkan bobot yang di-shard menjadi checkpoint reguler. Hal ini memerlukan accelerate>=1.12 agar Trainer dapat mencerminkan tp_size dan fsdp_size ke dalam ~Accelerate.ParallelismConfig.

Tabel di bawah ini membandingkan pelatihan EP-only dengan 2D EP+FSDP2 pada 8xH100 GPU. Beban kerja adalah fine-tuning penuh Qwen3-30B-A3B dalam bf16 pada panjang sekuens 2048. Semakin banyak shard FSDP, semakin rendah memori puncak, dan tokens/s turun sedikit karena FSDP2 melakukan all-gathers dan reduce-scatters para ahli di seluruh fsdp.

Melanjutkan dari checkpoint belum didukung untuk model yang di-shard saat dimuat, sehingga Trainer hanya menerima save_only_model=True atau save_strategy="no" untuk model tersebut.

Referensi API