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. Hầu hết các mô hình MoE mặc định dùng token dispatch, nên ep_size có thể được đặt độc lập với tp_size. Một số mô hình, chẳng hạn Llama 4 và Gemma 4, vẫn mặc định dùng masking và all-reduce ("ep_router" và "moe_tp_experts" trong model.ep_plan). Masking cũng khả dụng trên bất kỳ mô hình nào có ghi đè ep_plan, và yêu cầu ep_size=tp_size để mọi rank trong một nhóm chuyên gia nhận cùng các token.

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à ep_router che các chuyên gia nằm trên các rank khác trước khi all-reduce kết hợp các đầu ra chuyên gia.

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.

Quy tắc tiến của chuyên gia trong ep_plan xác định cách các token đến các chuyên gia:

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.

Hầu hết các mô hình MoE mặc định sử dụng "ep_dispatch_experts". Để sử dụng masking và all-reduce thay thế, hãy đặt ep_size=tp_size và ghi đè cả quy tắc forward của router lẫn expert (tên module router phụ thuộc vào mô hình, ví dụ mlp.router trên 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",
    },
)

Ngược lại, ghi đè quy tắc forward của expert trong một model có kế hoạch sử dụng masking bằng "ep_dispatch_experts" để dùng token dispatch. Khi đó, quy tắc router sẽ bị bỏ qua, vì dispatch cần các id expert toàn cục để tìm chủ sở hữu của từng expert.

Phân phối token

Với token dispatch, 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 các trọng số định tuyến. Chỉ các kích hoạt đượ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 tp, nhưng trạng thái optimizer và các module không có quy tắc vẫn được nhân bản (replicated) 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 mesh thứ hai với fsdp_size. Với masking và all-reduce, giữ ep_size=tp_size và truyền `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)

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