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. Model MoE memecah pakarnya dengan "ep_dispatch_experts". Implementasinya menentukan bagaimana token mencapai para pakar. Token dispatch memungkinkan ep_size berbeda dari tp_size. Beberapa implementasi EP mengharuskan ep_size=tp_size, misalnya "all_reduce" dan "deepgemm_megamoe", sehingga setiap rank dalam grup pakar memegang token yang sama.

Llama 4 tidak memiliki rencana EP: para ekspertnya mengambil setiap token yang sudah diskalakan oleh skor router-nya, sehingga tidak ada yang perlu didispatch. Pecah mereka dengan tensor parallelism (tp_size) sebagai gantinya.

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 tensor parallel memecah perhatian dan MLP padat, sedangkan rencana expert parallel 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 aturan "ep_dispatch_experts" pada modul ahli menentukan bagaimana token mencapai mereka.

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.

"ep_dispatch_experts" membagi seluruh ahli per rank (ep_size adalah kelipatan dari tp_size yang membagi fsdp_size * tp_size) dan mengaksesnya melalui implementasi yang dinamakan oleh ep_implementation, lihat Memilih implementasi.

Memilih implementasi

"ep_dispatch_experts" memecah para ahli sekali dan memilih cara token mencapai mereka dari implementasi yang terdaftar di EPImplementationInterface, yang dipilih dengan ep_implementation di DistributedConfig:

ep_implementation hanya memilih di antara ukuran paralel yang diatur dalam DistributedConfig. Ketika rank dari sebuah grup ahli memegang batch yang berbeda, misalnya tp_size=1 dengan fsdp_size=ep_size, "all_to_all" adalah satu-satunya implementasi yang berlaku.

Menambahkan implementasi

Implementasi mewarisi EPImplementation dan terdaftar di bawah nama yang dipilih oleh ep_implementation. Ia dibangun sekali per modul pakar, sehingga dapat menyimpan state seperti buffer komunikasi. Memanggilnya dengan batch yang dipegang oleh rank ini, beserta id pakar global, mengembalikan output pakar berbobot untuk masing-masing tokennya. Ia menjalankan pakar milik rank ini dengan experts_forward, pada baris, id pakar lokal, dan bobot yang dibawanya.

from transformers.distributed.expert_parallel import EPImplementation, EPImplementationInterface

class MyEP(EPImplementation):
    def __call__(self, hidden_states, top_k_index, top_k_weights):
        tokens, expert_ids, weights = ...  # bring this rank's experts the rows routed to them
        expert_output = self.experts_forward(tokens, expert_ids, weights)
        return ...  # the weighted expert output for each token of the batch this rank holds

EPImplementationInterface.register("my_ep", MyEP)
distributed_config = DistributedConfig(tp_size=1, fsdp_size=8, ep_size=8, ep_implementation="my_ep")

Gradien harus mencapai setiap rank yang memegang token, termasuk bobot routing, sehingga bungkus kolektif yang tidak didiferensiasi oleh autograd dalam torch.autograd.Function.

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.

Dengan ep_size=tp_size, implementasi default "auto" menjalankan "all_reduce" saat inferensi dan mendistribusikan token saat pelatihan. Untuk melakukan masking dan all-reduce juga saat pelatihan, atur ep_implementation="all_reduce":

distributed_config = DistributedConfig(tp_size=4, ep_size=4, ep_implementation="all_reduce")

Token dispatch

Token dispatch adalah implementasi "all_to_all": setiap rank melatih pada bagian batch-nya sendiri. Di setiap lapisan MoE, sebuah rank merutekan tokennya, mengirim setiap pasangan (token, expert) ke rank yang memiliki expert tersebut dengan all-to-all, menjalankan expert lokalnya pada apa yang diterimanya, mendapatkan hasilnya kembali dengan 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 membagi (shard) bobot di atas mesh mereka sendiri, tp dan ep, tetapi state optimizer dan modul tanpa aturan masih direplikasi di setiap rank dari grup tersebut, yang membatasi seberapa besar model yang dapat dilatih. Tambahkan FSDP2 pada dimensi mesh lain dengan fsdp_size. Dengan masking dan all-reduce, pertahankan ep_size=tp_size dan lewati ep_implementation="all_reduce".

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_implementation="all_reduce",  # "auto" would dispatch tokens in training, which the Trainer does not support yet
)
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