Expert Parallelism

Song song chuyên gia (EP) là chiến lược huấn luyện phân tán và suy luận cho các model hỗn hợp chuyên gia (MoE), trong đó các chuyên gia khác nhau được đặt trên các GPU riêng biệt và mỗi token được gửi đến thiết bị chứa các chuyên gia được chọn cho token đó. Một lớp MoE thay khối truyền thẳng lớn của Transformer bằng nhiều mạng chuyên gia nhỏ hơn cùng bộ định tuyến chỉ kích hoạt một vài chuyên gia cho mỗi token; do đó, tổng số tham số có thể tăng nhanh hơn nhiều so với lượng tính toán cho mỗi token. Song song chuyên gia giúp các tham số đó vừa trong bộ nhớ: thay vì mỗi GPU lưu mọi chuyên gia, mỗi GPU chỉ lưu một phần. Từ năm 2017, Noam Shazeer cùng các đồng nghiệp tại Google đã phân tán chuyên gia trên các thiết bị, mỗi chuyên gia có một bản dùng chung, trong lớp MoE điều khiển thưa của họ; đến tháng 6 năm 2020, GShard mở rộng phương pháp này bằng cách phân tán các chuyên gia của model dịch thuật hơn 600 tỷ tham số trên 2.048 bộ tăng tốc TPU v3. Song song chuyên gia hiện là tiêu chuẩn để huấn luyện và phục vụ các model ngôn ngữ và thị giác-ngôn ngữ MoE quy mô lớn.

Token di chuyển giữa các GPU như thế nào#

Mỗi layer MoE hoạt động theo bốn bước:

1. Định tuyến: Bộ định tuyến chấm điểm từng token với từng chuyên gia và chọn các chuyên gia top-k, thường là một, hai hoặc tám.
2. Điều phối: Một phép toán tập thể all-to-all gửi trạng thái ẩn của từng token đến các GPU sở hữu những chuyên gia được chọn.
3. Tính toán: Mỗi GPU chạy các chuyên gia cục bộ trên những token mà GPU đó nhận được.
4. Kết hợp: Một phép all-to-all thứ hai trả kết quả về GPU ban đầu của từng token; tại đó, kết quả được tính trọng số theo điểm của bộ định tuyến rồi cộng lại.

Đoạn mã dưới đây mô phỏng các bước định tuyến và phân phối cho một layer có 8 chuyên gia trải trên 4 GPU, mỗi GPU có hai chuyên gia:

import torch

torch.manual_seed(0)
num_experts, num_gpus, top_k = 8, 4, 2
tokens = torch.randn(4096, 512)  # trạng thái ẩn cho 4.096 token
router = torch.nn.Linear(512, num_experts)

# Định tuyến: mỗi token chọn 2 chuyên gia hàng đầu theo điểm của bộ định tuyến
with torch.no_grad():
    _, experts = router(tokens).topk(top_k, dim=-1)

# Song song chuyên gia: chuyên gia 0-1 nằm trên GPU 0, chuyên gia 2-3 trên GPU 1, v.v.
gpu = experts // (num_experts // num_gpus)

# Phân phối: đếm số lượt gán token-chuyên gia mà mỗi chuyên gia và GPU nhận được
print(torch.bincount(experts.flatten(), minlength=num_experts).tolist())
print(torch.bincount(gpu.flatten(), minlength=num_gpus).tolist())

Dòng đầu tiên trong kết quả cho biết tải của mỗi chuyên gia, từ 948 đến 1.087 lượt gán, so với mức chia đều là 1.024; dòng thứ hai cho biết bốn GPU nhận từ 1.968 đến 2.139 trong tổng số 8.192 lượt gán. Ngay cả bộ định tuyến được khởi tạo ngẫu nhiên cũng phân bổ công việc không đồng đều; bộ định tuyến đã huấn luyện còn hình thành các chuyên gia được ưa chuộng, khiến tải lệch nhiều hơn.

Cân bằng tải và dung lượng#

Một layer chỉ hoàn tất khi GPU bận nhất hoàn tất, vì vậy mất cân bằng định tuyến trực tiếp khiến phần cứng nhàn rỗi. Huấn luyện MoE khắc phục vấn đề này theo ba cách:

• Loss cân bằng phụ trợ: Một số hạng bổ sung trong hàm loss khuyến khích bộ định tuyến phân bổ token đồng đều giữa các chuyên gia.
• Hệ số dung lượng: Mỗi chuyên gia chỉ nhận tối đa một số token cố định trong mỗi batch, chẳng hạn 1.25 lần phần được phân bổ đồng đều; token vượt giới hạn sẽ bỏ qua chuyên gia đó.
• Cân bằng dựa trên bias: DeepSeek-V3 cộng bias riêng cho từng chuyên gia vào điểm số của bộ định tuyến; bias này chỉ ảnh hưởng đến việc chọn chuyên gia, giảm đối với chuyên gia quá tải và tăng đối với chuyên gia ít hoạt động sau mỗi bước. Cách này cân bằng tải mà không cần loss phụ trợ cạnh tranh với mục tiêu chính.

Khi có giới hạn dung lượng, đường đi của mỗi token phụ thuộc vào mức độ bận của các chuyên gia mà token đó được định tuyến đến:

Token bị bỏ qua tiếp tục đi qua lớp bằng kết nối tắt mà không thay đổi, vì vậy loại bỏ quá nhiều token sẽ làm giảm chất lượng. DeepSeek-V3 không loại bỏ token nào trong quá trình huấn luyện mà thay vào đó dựa vào cơ chế cân bằng bằng bias.

Chi phí còn lại là truyền thông. Hai lần trao đổi all-to-all diễn ra trong mọi lớp MoE, vì vậy song song chuyên gia hoạt động tốt nhất bên trong một node, nơi các GPU dùng chung kết nối NVLink tốc độ cao; truyền thông qua các liên kết chậm hơn giữa các node sẽ tốn kém. DeepSeek-V3 phân tán 256 chuyên gia được định tuyến trong mỗi lớp bằng song song chuyên gia 64 chiều trên 8 node, gửi mỗi token đến các chuyên gia trên tối đa 4 node và chồng lấp quá trình trao đổi với tính toán.

Kết hợp các chiến lược song song#

Song song chuyên gia là một trong nhiều cách phân chia quá trình huấn luyện giữa các thiết bị:

• Song song dữ liệu: Nhân bản toàn bộ model trên mỗi GPU và chia mỗi batch, đồng bộ hóa gradient sau mỗi bước.
• Song song tensor: Chia các ma trận trọng số riêng lẻ giữa các GPU, khiến mọi GPU tham gia vào từng lớp.
• Song song pipeline: Phân các lớp liên tiếp cho những GPU khác nhau rồi truyền các micro-batch lần lượt qua chúng.
• Song song ngữ cảnh: Chia một chuỗi đầu vào dài giữa các GPU.

Các model dense vừa vặn trên một GPU không cần gì ngoài data parallelism. Đây là cách Ultralytics YOLO26 huấn luyện trên nhiều GPU khi truyền danh sách thiết bị như device=[0, 1] trong training mode. Các model MoE lớn kết hợp nhiều phương pháp: DeepSeek-V3 đã huấn luyện 671 tỷ tham số, trong đó có 37 tỷ tham số hoạt động trên mỗi token, bằng pipeline parallelism 16 chiều, expert parallelism 64 chiều và data parallelism phân mảnh, không dùng tensor parallelism. Các inference engine như vLLM áp dụng cách phân chia tương tự khi phục vụ một model, nhờ đó model MoE có các expert vượt quá bộ nhớ của một GPU vẫn có thể chạy trên nhiều GPU, mỗi GPU chỉ lưu phần expert được phân bổ.