Uzman paralelliği

Uzman paralelliği, mixture-of-experts (MoE) modelleri için bir paralellik stratejisidir. Her uzmanın ileri besleme katmanı farklı bir donanım hızlandırıcısında bulunur. Bir yönlendirici, token'ları uygun uzmanlara yönlendirir ve sonuçları toplar. Her token yalnızca birkaç uzmanı etkinleştirdiği için bu yaklaşım, hesaplama maliyetini artırmadan modelleri çok daha büyük parametre sayılarına ölçeklendirir.

DistributedConfig

DistributedConfig sınıfını ve ep_size argümanını kullanarak uzman paralelliğini etkinleştirin. MoE modelleri uzmanlarını "ep_dispatch_experts" ile bölümler. Uygulamaları, tokenlerin uzmanlara nasıl ulaştığını belirler. Token dağıtımı, ep_size'in tp_size'ten farklı olmasına olanak tanır. Bazı EP uygulamaları ep_size=tp_size gerektirir; örneğin "all_reduce" ve "deepgemm_megamoe", böylece bir uzman grubundaki her rank aynı tokenlere sahip olur.

Llama 4'te EP planı yoktur: uzmanları, yönlendirici puanıyla zaten ölçeklenmiş her token'ı alır, bu nedenle dağıtılacak bir şey yoktur. Bunun yerine onları tensör paralelliği (tp_size) ile bölümlendirin.

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

Her MoE modeli, config dosyasında iki plan tanımlar: yoğun modüller için base_model_tp_plan ve uzmanlar için base_model_ep_plan. Bunlar yüklenen modelde model.tp_plan ve model.ep_plan olarak erişilebilir hale gelir. tp_size > 1 ve ep_size > 1 olduğunda her ikisi de uygulanır: tensor parallel planı dikkat (attention) ve yoğun MLP'leri parçalar (shard), uzman paralel planı ise uzmanları parçalar. Aynı modüller için EP kuralları TP kurallarına öncelik taşır, bu nedenle uzman ağırlıkları yalnızca EP planı tarafından bir kez parçalanır. EP planında, GroupedGemmParallel stili uzman ağırlıklarını uzman boyutu boyunca böler, böylece her rank yalnızca yerel uzmanlarını yükler; uzman modülündeki "ep_dispatch_experts" kuralı ise token'ların bu uzmanlara nasıl ulaştığını belirler.

tp_plan yalnızca tp_size > 1 olduğunda, ep_plan ise yalnızca ep_size > 1 olduğunda uygulanır. TP etkin ve EP devre dışı olduğunda, uzman kuralları dahil olmak üzere tam TP planı uygulanır.

"ep_dispatch_experts", her rank için tüm uzmanları parçalar (ep_size, tp_size'nin bir katı olup fsdp_size * tp_size'i bölen bir değerdir) ve bunlara ep_implementation adını taşıyan uygulama üzerinden erişir, bkz. Uygulama seçimi.

Bir uygulama seçimi

"ep_dispatch_experts", uzmanları bir kez parçalar ve DistributedConfig içindeki ep_implementation ile seçilen, EPImplementationInterface’te kayıtlı uygulamalardan tokenlerin onlara nasıl ulaşacağını belirler:

ep_implementation, yalnızca DistributedConfig içinde belirlenen paralel boyutlar arasından seçim yapar. Bir uzman grubundaki rank'lar farklı partileri tuttuğunda, örneğin tp_size=1 ve fsdp_size=ep_size durumunda, uygulanabilir tek uygulama "all_to_all"'dır.

Bir uygulama ekleme

Bir uygulama EPImplementation sınıfından türetilir ve ep_implementation adıyla seçilen kayıt altına alınır. Uzman modülü başına bir kez oluşturulduğundan, iletişim tamponları gibi durumları koruyabilir. Bu rank'ın tuttuğu batch ile, küresel uzman kimlikleriyle çağrıldığında, token'larının her biri için ağırlıklı uzman çıktısını döndürür. Bu rank'ın uzmanlarını, getirdiği satırlar, yerel uzman kimlikleri ve ağırlıklar üzerinde experts_forward ile çalıştırır.

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")

Gradyanlar, token tutan her ranga, yönlendirme ağırlıkları dahil olmak üzere, ulaşmalıdır; bu nedenle, autograd'ın türevini almadığı bir kolektif işlemi torch.autograd.Function içinde sarmalayın.

enable_expert_parallel=True, ep_size olmadan kullanıldığında artık desteklenmemektedir ve bir FutureWarning üretir. Bu ayar, yalnızca uzmanları parçalar; tıpkı ep_size=tp_size, tp_plan={} yapılandırmasında olduğu gibi.

Çıkarım betiğinizi torchrun ile başlatın. İşlem sayısı tp_size * fsdp_size * pp_size değerine eşit olmalı ve ep_size, uzman sayısını tam olarak bölmelidir.

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

Planların üzerine yazılması

tp_plan={...} veya ep_plan={...} değerlerini DistributedConfig sınıfına geçirerek önceden tanımlanmış planların bireysel kurallarını geçersiz kılabilirsiniz. Belirtilmeyen kurallar korunur ve birleştirilen planlar modele kaydedilir. Her anahtar, bir modül, bir parametre veya mevcut bir plan girdisiyle eşleşmelidir; aksi takdirde herhangi bir şey bölümlenmeden önce yükleme işlemi ValueError hatası verir. Tam yolu, yüklü modelin gördüğü şekilde kullanın; bu nedenle bir nedensel LM için model.layers.*, temel modeli için ise layers.* kullanın.

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"},
)

Bir plan sağlanmak, paralel boyutları çıkarım yapmaz: tp_size ve ep_size değerlerini açıkça belirleyin.

ep_size=tp_size durumunda, varsayılan "auto" uygulaması çıkarım sırasında "all_reduce" çalıştırır ve eğitim sırasında jetonları dağıtır. Eğitimde de maskeleme ve all-reduce yapmak için ep_implementation="all_reduce" ayarlayın:

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

Token dağıtımı

Token dispatch, "all_to_all" uygulamasıdır: her rank, batch'in kendi payına düşen kısmında eğitim yapar. Her MoE katmanında, bir rank token'larını yönlendirir, her (token, uzman) çiftini, o uzmanın sahibi olduğu rank'a bir all-to-all ile gönderir, yerel uzmanlarını aldıkları üzerinde çalıştırır, sonuçları ikinci bir all-to-all ile geri alır ve bunları yönlendirme ağırlıklarıyla birleştirir. Yalnızca yönlendirilen aktivasyonlar ve uzman çıktıları taşınır ve hiçbir rank, sahip olmadığı token'lar için uzman hesaplama yapmaz.

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)

tp_size=1 durumunda, ep_size hem fsdp_size'in hem de uzman sayısının tam bölünmesini sağlamalıdır ve dikkat mekanizması num_key_value_heads ile sınırlı değildir. Modelin geri kalanı için:

• Uzmanların dışındaki parametreler, fsdp üzerinde FSDP2 ile parçalanır ve bu da onların gradyanlarını azaltır.
• Uzmanlar ep üzerinde parçalanır ve efsdp_size = fsdp_size * tp_size // ep_size birinden büyük olduğunda, ayrıca efsdp üzerinde FSDP ile parçalanır. Her zaman FSDP ile sarılı olduklarından, fsdp_mixed_precision ve fsdp_cpu_offload onlara da uygulanır ve save_pretrained() onları diğer herhangi bir parametre gibi toplar.
• Uzman paralel bir grup ep_size / tp_size parti tutar, bu nedenle bir uzmanın gradyanı bu kadar partinin toplamıdır. efsdp indirgeme işlemi, grup boyutu yerine fsdp_size değerine böler; bu da FSDP2'nin yoğun modüller için kullandığı parti başına ortalama ile aynı sonucu verir.
• Her yerel uzman, katman başına bir sıfır dolgu satırını da işler. Uzmanları tarafından hiçbir token alınmayan bir sıralama, yine de reverse all-to-all ve uzman gradyan indirgeme işlemlerine katılır.

Boyutların nasıl birleştiği

• tp_size * fsdp_size, işlem sayısıdır. ep_size hiçbir şey eklemez: yalnızca uzman ağırlıkları için aynı sıraları yeniden gruplar.
• ep_size, uzman listesini ep_size bloğa böler. Her rank num_experts / ep_size uzman hesaplar ve ep_size ardışık rank bir tam seti barındırır. Bu rank grubu, all-to-all'ın çalıştığı gruptur.
• efsdp_size = fsdp_size * tp_size / ep_size, uzman kümesinin kaç tam kopyasının var olduğunu belirtir. Farklı kopyalarda aynı konumdaki rank'lar, FSDP'nin yoğun modüller için yaptığı gibi, bu uzmanları bellek için parçalara ayırır ve gradyanlarını ortalama alır.
• Bir rank'ın tuttuğu batch yalnızca fsdp'ye bağlıdır. Ardışık rank'lar bir TP grubu oluşturur ve aynı batch'i alır; fsdp_size grupları farklı batch'ler alır.

Sekiz süreç, tp_size=2, fsdp_size=4, sekiz uzman:

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

Şekilden iki sonuç çıkar:

• Bir EP grubunun içinde, her token, kendi batch'ini tutan çiftin rank'ı başına bir kez olmak üzere tp_size kez var olur. Bu, ep_size'dan bağımsızdır.
• Bir EP grubu ep_size / tp_size farklı batch tutar. Bu, bir uzmanın gradyanlarının üzerinde toplandığı sayıdır.

Tensor paralelliği ile

Yoğun modülleri TP planıyla parçalara ayırmak için tp_size > 1 değerini ayarlayın; uzmanlar ise dispatch kullanır. Sekiz işlemde:

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

Her TP rank çifti aynı batch’i alır, çünkü tensor paralelliği çiftin içindeki aktivasyonları kopyalar. Her iki rank da tüm token’larını dağıttıysa, bir uzmanı (expert) barındıran rank her token’ı iki kez alır, iki kez hesaplar ve ağırlık gradyanı ikiye katlanır. Uzmanlar tek bir rank üzerinde bütündür, bu nedenle kopya ağırlıklar üzerinden bölünemez ve sahibi genellikle başka bir rank olduğu için, masking’de olduğu gibi sahiplik üzerinden de çözülemez. Çift bu nedenle satırları böler: her TP rank, token’ların ayrık 1 / tp_size’ını dağıtır, sonuçlar gönderen rank’a geri döner ve sıfır dolgulu yarıların çift üzerinde all-reduce’ı, bir sonraki katmanın beklediği kopyalanmış çıktıyı geri getirir. Bölme tp_size üzerinden yapılır, ep_size üzerinden değil, çünkü yalnızca bir batch tutan rank’lar onu gönderebilir. Uzman grupları dört rank’ı kapsar ve her uzman, efsdp_size = 2 rank arasında FSDP ile bölünürken, gövdenin (trunk) FSDP grubu dört rank’ı kapsar. Modelin olağan TP kısıtlamaları, örneğin dikkat başlarının bölünebilirliği, yoğun (dense) modüller için hâlâ geçerlidir. Token dilimleri eşitsiz veya boş olabilir, tek token çözümlendirme sırasında da dahil.

Token dispatch henüz pipeline paralelliğiyle birleştirilemez; pp_size=1 kullanın (henüz test edilmedi) Bu yapılandırmaların her biri sekiz GPU kullanır:

Trainer henüz token dağıtımını hesaba katmıyor: batch ve token sayımı, bir TP grubundaki rank'lerin bir batch'i paylaşması ve fsdp_size adet veri paralel diliminin var olduğu all-reduce yerleşimini varsayar. Trainer'ın dağıtıma desteği bir sonraki aşamada eklenecek.

FSDP2 ile birleştirme

Tensor ve uzman paralelliği, ağırlıkları kendi mesh'leri olan tp ve ep üzerinde parçalara ayırır; ancak optimizör durumu ve bir kuralı olmayan modüller, grubun her rank'ında hâlâ replike edilir ve bu da eğitebileceğiniz modelin boyutunu sınırlar. fsdp_size ile başka bir mesh boyutunda FSDP2 ekleyin. Masking ve all-reduce ile ep_size=tp_size değerini koruyun ve ep_implementation="all_reduce" parametresini iletin.

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)

Modeli normal şekilde yükleyin, ardından Trainer ile eğitim yapın. Gradyan normunu her iki mesh üzerinden alır ve her mesh'e kendi optimizer param grubunu verir. save_model(), parçalı ağırlıkları düzenli bir checkpoint'e toplar. Bunun için Trainer'ın tp_size ve fsdp_size değerlerini ~Accelerate.ParallelismConfig'e yansıtabilmesi gerektirdiğinden accelerate>=1.12 sürümü gereklidir.

Aşağıdaki tablo, 8xH100 GPU'larda yalnızca EP ile eğitimi, 2B EP+FSDP2 ile karşılaştırır. İş yükü, 2048 dizi uzunluğunda bf16'da Qwen3-30B-A3B'nin tam ince ayarıdır. Daha fazla FSDP parçası tepe belleği keser ve FSDP2 uzmanları fsdp üzerinden all-gather ve reduce-scatter yaptığı için token/s biraz düşer.

Yükleme sırasında parçalanan modeller için kontrol noktasından devam etme henüz desteklenmediğinden, Trainer bu modeller için yalnızca save_only_model=True veya save_strategy="no" değerlerini kabul eder.

API referansı