Song song chuyên gia

Expert parallelism là một chiến lược song song hóa cho các mô hình mixture-of-experts (MoE). Lớp feedforward của mỗi expert nằm trên một bộ tăng tốc phần cứng khác nhau. Một router sẽ điều phối các token đến các expert phù hợp và thu thập kết quả. Cách tiếp cận này cho phép mở rộng mô hình lên số lượng tham số lớn hơn nhiều mà không làm tăng chi phí tính toán, vì mỗi token chỉ kích hoạt một vài expert.

DistributedConfig

Bật song song chuyên gia bằng lớp DistributedConfig và đối số ep_size. Các mô hình MoE chia nhỏ chuyên gia của chúng bằng "ep_dispatch_experts". Các cài đặt của nó quyết định cách các token đến các chuyên gia. Token dispatch cho phép ep_size khác với tp_size. Một số cài đặt EP yêu cầu ep_size=tp_size, ví dụ "all_reduce" và "deepgemm_megamoe", để mọi rank trong một nhóm chuyên gia giữ cùng các token.

Llama 4 không có kế hoạch EP: các chuyên gia của nó nhận mọi token đã được nhân theo điểm router, nên không có gì để điều phối. Hãy chia nhỏ chúng bằng song song tensor (tp_size) thay vào đó.

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

Mỗi mô hình MoE định nghĩa hai kế hoạch trong cấu hình của nó: base_model_tp_plan cho các module dense và base_model_ep_plan cho các chuyên gia. Chúng được hiển thị trên mô hình đã tải dưới dạng model.tp_plan và model.ep_plan. Với tp_size > 1 và ep_size > 1, cả hai đều được áp dụng: kế hoạch tensor parallel chia nhỏ attention và các MLP dense, còn kế hoạch expert parallel chia nhỏ các chuyên gia. Các quy tắc EP có ưu tiên cao hơn các quy tắc TP đối với cùng một module, do đó trọng số chuyên gia chỉ được chia nhỏ một lần, bởi kế hoạch EP. Trong kế hoạch EP, kiểu GroupedGemmParallel chia nhỏ trọng số chuyên gia dọc theo chiều chuyên gia để mỗi rank chỉ tải các chuyên gia cục bộ của nó, và quy tắc "ep_dispatch_experts" trên module chuyên gia xác định cách các token được chuyển đến chúng.

tp_plan chỉ được áp dụng khi tp_size > 1, và ep_plan chỉ khi ep_size > 1. Khi TP được bật và EP bị tắt, toàn bộ kế hoạch TP sẽ được áp dụng, bao gồm cả các quy tắc chuyên gia.

"ep_dispatch_experts" chia các chuyên gia nguyên vẹn theo từng rank (ep_size là bội số của tp_size và chia hết fsdp_size * tp_size) và truy cập chúng thông qua cài đặt do ep_implementation chỉ định, xem Chọn một cài đặt.

Chọn một triển khai

"ep_dispatch_experts" chia nhỏ các chuyên gia một lần và chọn cách các token tiếp cận chúng từ các cài đặt được đăng ký trong EPImplementationInterface, được chọn bằng ep_implementation trong DistributedConfig:

ep_implementation chỉ chọn trong số các kích thước song song được đặt trong DistributedConfig. Khi các rank của một nhóm chuyên gia giữ các batch khác nhau, ví dụ tp_size=1 với fsdp_size=ep_size, "all_to_all" là triển khai duy nhất được áp dụng.

Thêm một cài đặt

Một implementation kế thừa EPImplementation và đăng ký dưới tên ep_implementation được chọn. Nó được xây dựng một lần cho mỗi module experts, do đó có thể giữ trạng thái như các bộ đệm truyền thông. Khi được gọi với batch mà rank này đang giữ, cùng với các expert id toàn cục, nó trả về đầu ra expert đã được trọng số hóa cho từng token của nó. Nó chạy các experts của rank này bằng experts_forward, trên các hàng, expert id cục bộ và trọng số mà nó đã mang đến cho chúng.

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

Gradient phải đến được mọi rank đang giữ một token, bao gồm cả trọng số định tuyến, vì vậy hãy bọc một phép toán tập thể mà autograd không phân biệt trong một torch.autograd.Function.

enable_expert_parallel=True mà không có ep_size đã bị ngừng hỗ trợ và sẽ phát ra FutureWarning. Nó chỉ chia nhỏ các chuyên gia, giống như ep_size=tp_size, tp_plan={}.

Chạy script suy luận của bạn bằng torchrun. Số lượng tiến trình phải bằng tp_size * fsdp_size * pp_size, và ep_size phải chia hết cho số lượng chuyên gia.

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

Ghi đè các kế hoạch

Truyền tp_plan={...} hoặc ep_plan={...} vào DistributedConfig để ghi đè các quy tắc riêng lẻ của các kế hoạch được định nghĩa trước. Các quy tắc không được chỉ định sẽ được giữ nguyên, và các kế hoạch đã hợp nhất sẽ được lưu trên model. Mỗi key phải khớp với một module, một parameter, hoặc một mục có sẵn trong kế hoạch; nếu không, quá trình tải sẽ gây ra ValueError trước khi bất kỳ phần nào được sharding. Hãy sử dụng đường dẫn đầy đủ như nhìn từ model đã được tải, ví dụ model.layers.* cho một causal LM và layers.* cho base model của 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"},
)

Việc cung cấp một kế hoạch không suy ra kích thước song song: hãy đặt tp_size và ep_size một cách rõ ràng.

Với ep_size=tp_size, cài đặt mặc định "auto" chạy "all_reduce" ở giai đoạn suy luận và phân phối token trong huấn luyện. Để che (mask) và all-reduce cả trong huấn luyện, hãy đặt ep_implementation="all_reduce":

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

Phân phối token

Token dispatch là phần triển khai "all_to_all": mỗi rank huấn luyện trên phần riêng của batch. Tại mỗi lớp MoE, một rank định tuyến các token của nó, gửi mỗi cặp (token, expert) đến rank sở hữu expert đó bằng all-to-all, chạy các expert cục bộ trên những gì nó nhận được, nhận kết quả trở lại bằng một all-to-all thứ hai và kết hợp chúng với trọng số định tuyến. Chỉ các activation được định tuyến và đầu ra của expert di chuyển, và không rank nào tính toán expert cho các token mà nó không sở hữu.

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)

Với tp_size=1, ep_size phải chia hết cho fsdp_size và số lượng chuyên gia, và attention không bị ràng buộc bởi num_key_value_heads. Đối với phần còn lại của mô hình:

• Các tham số nằm ngoài các chuyên gia được chia nhỏ bằng FSDP2 trên fsdp, điều này làm giảm gradient của chúng.
• Các chuyên gia được chia nhỏ (sharded) trên ep và, khi efsdp_size = fsdp_size * tp_size // ep_size lớn hơn một, chúng còn được chia nhỏ theo FSDP trên efsdp. Chúng luôn được bọc (wrapped) bởi FSDP, do đó fsdp_mixed_precision và fsdp_cpu_offload cũng áp dụng cho chúng và save_pretrained() sẽ thu thập (gather) chúng giống như bất kỳ tham số nào khác.
• Một nhóm song song chuyên gia giữ ep_size / tp_size lô, do đó gradient của một chuyên gia là tổng trên số lô đó. Việc giảm efsdp chia cho fsdp_size thay vì kích thước nhóm của nó, cho ra cùng một giá trị trung bình trên mỗi lô mà FSDP2 áp dụng cho các module đặc.
• Mỗi chuyên gia cục bộ cũng xử lý một hàng đệm zero cho mỗi lớp. Một rank mà các chuyên gia của nó không nhận được token nào vẫn tham gia vào quá trình reverse all-to-all và giảm gradient chuyên gia.

Cách các kích thước kết hợp

• tp_size * fsdp_size là số lượng tiến trình. ep_size không thêm gì: nó chỉ nhóm lại các rank giống nhau cho trọng số chuyên gia.
• ep_size chia danh sách chuyên gia thành ep_size khối. Mỗi rank tính toán num_experts / ep_size chuyên gia, và ep_size rank liên tiếp giữ một bộ đầy đủ. Tập hợp các rank đó là nhóm mà all-to-all chạy trong đó.
• efsdp_size = fsdp_size * tp_size / ep_size là số lượng bản sao đầy đủ của tập chuyên gia tồn tại. Các rank ở cùng vị trí trong các bản sao khác nhau sẽ chia nhỏ (shard) các chuyên gia đó để tiết kiệm bộ nhớ và trung bình hóa gradient của chúng, giống như cách FSDP thực hiện cho các module dense.
• Batch mà một rank giữ phụ thuộc duy nhất vào fsdp. Các rank liên tiếp tạo thành một nhóm TP và nhận cùng một batch; fsdp_size nhóm nhận các batch khác nhau.

Tám tiến trình, tp_size=2, fsdp_size=4, tám chuyên gia:

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

Hai con số rút ra từ hình ảnh:

• Bên trong một nhóm EP, mỗi token tồn tại tp_size lần, một lần cho mỗi rank của cặp giữ batch của nó. Điều này không phụ thuộc vào ep_size.
• Một nhóm EP giữ ep_size / tp_size các lô dữ liệu khác nhau. Đây là số lượng mà gradient của một chuyên gia được cộng dồn.

Với song song tensor

Đặt tp_size > 1 để chia nhỏ các module dense theo kế hoạch TP trong khi các expert sử dụng dispatch. Trên tám tiến trình:

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

Mỗi cặp rank TP nhận cùng một batch, vì tensor parallelism nhân bản các activation bên trong cặp đó. Nếu cả hai rank đều dispatch toàn bộ token của mình, rank sở hữu một expert sẽ nhận mỗi token hai lần, tính toán hai lần, và gradient trọng số của nó sẽ bị nhân đôi. Các expert nằm trọn trên một rank, nên bản sao thừa không thể được chia theo trọng số, và rank sở hữu thường là một rank khác, nên cũng không thể giải quyết bằng quyền sở hữu như cách masking làm. Do đó, cặp này chia các hàng: mỗi rank TP dispatch một phần rời rạc 1 / tp_size của các token, kết quả được trả về rank đã gửi, và một all-reduce trên cặp của các nửa được zero-pad khôi phục đầu ra được nhân bản mà lớp tiếp theo kỳ vọng. Việc chia được thực hiện theo tp_size, không phải ep_size, vì chỉ các rank giữ một batch mới có thể gửi nó. Các nhóm expert trải dài trên bốn rank và mỗi expert được FSDP-sharded trên efsdp_size = 2 rank, trong khi nhóm FSDP của trunk trải dài trên bốn rank. Các ràng buộc TP thông thường của mô hình, chẳng hạn như tính chia hết của attention-head, vẫn áp dụng cho các module dense. Các lát token có thể không đều hoặc rỗng, kể cả trong quá trình giải mã đơn token.

Token dispatch hiện chưa thể kết hợp với pipeline parallelism; hãy sử dụng pp_size=1 (chưa được kiểm tra) Các cấu hình này đều sử dụng tám GPU:

Trainer hiện chưa hỗ trợ việc phân phối token: việc đếm batch và token giả định bố cục all-reduce, trong đó các rank của một nhóm TP chia sẻ một batch và tồn tại fsdp_size shard song song dữ liệu. Hỗ trợ dispatch cho Trainer sẽ được bổ sung trong một bản cập nhật tiếp theo.

Kết hợp với FSDP2

Tensor và expert parallelism chia nhỏ (shard) trọng số trên các lưới riêng của chúng, tp và ep, nhưng trạng thái optimizer và các module không có quy tắc vẫn được nhân bản (replicate) trên mọi rank của nhóm, điều này giới hạn kích thước mô hình mà bạn có thể huấn luyện. Thêm FSDP2 trên một chiều lưới khác với fsdp_size. Với masking và all-reduce, giữ ep_size=tp_size và truyền 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)

Tải model như bình thường, sau đó huấn luyện với Trainer. Nó lấy chuẩn gradient trên cả hai mesh và cấp cho mỗi mesh một nhóm tham số optimizer riêng. save_model() thu thập các trọng số đã chia mảnh thành một checkpoint thông thường. Điều này yêu cầu accelerate>=1.12 để Trainer có thể sao chép tp_size và fsdp_size vào ~Accelerate.ParallelismConfig.

Bảng dưới đây so sánh huấn luyện chỉ dùng EP với 2D EP+FSDP2 trên 8 GPU H100. Khối lượng công việc là fine-tuning đầy đủ Qwen3-30B-A3B ở bf16 với độ dài chuỗi 2048. Nhiều shard FSDP hơn sẽ cắt giảm bộ nhớ đỉnh, và tokens/s giảm một chút vì FSDP2 thực hiện all-gather và reduce-scatter các expert qua fsdp.

Việc tiếp tục từ một checkpoint hiện chưa được hỗ trợ cho các mô hình được chia mảnh (sharded) tại thời điểm tải, do đó Trainer chỉ chấp nhận save_only_model=True hoặc save_strategy="no" cho chúng.

Tham khảo API