Hỗn hợp Chuyên gia (MoE) trong Transformers
Tối ưu hóa MoE trong Transformers
- 12 min read
Mixture of Experts (MoEs) in Transformers
Introduction
Trong những năm gần đây, việc mở rộng quy mô các mô hình ngôn ngữ lớn (LLMs) dày đặc đã thúc đẩy phần lớn tiến bộ trong LLMs. Từ các mô hình ban đầu như ULMFiT (~30 triệu tham số) hoặc GPT-2 (1.5 tỷ tham số, từng bị coi là “quá nguy hiểm để phát hành” 🧌), cho đến các hệ thống hàng trăm tỷ tham số ngày nay, công thức rất đơn giản:
Nhiều dữ liệu + nhiều tham số mang lại hiệu suất tốt hơn.
Các định luật về quy mô (Scaling laws) đã củng cố xu hướng này, nhưng việc mở rộng quy mô dày đặc có những giới hạn thực tế:
- Việc đào tạo ngày càng tốn kém.
- Độ trễ suy luận tăng lên.
- Việc triển khai đòi hỏi bộ nhớ và phần cứng đáng kể.
Đây là nơi Mixture of Experts (MoEs) xuất hiện.
From Dense to Sparse: What Are MoEs?
Mô hình Mixture of Experts giữ lại cấu trúc Transformer cơ bản, nhưng thay thế các lớp feed-forward dày đặc bằng một tập hợp các chuyên gia (experts). “Chuyên gia” không phải là một mô-đun chuyên biệt theo chủ đề (ví dụ: “chuyên gia toán học”, “chuyên gia mã”). Nó chỉ đơn giản là một mạng con có thể học được. Đối với mỗi token, một bộ định tuyến (router) sẽ chọn một tập hợp nhỏ các chuyên gia để xử lý nó.
Các token khác nhau kích hoạt các chuyên gia khác nhau, dựa trên biểu diễn ẩn của chúng.
Đây là ý tưởng cốt lõi.
Ví dụ, hãy xem xét gpt-oss-20b. Nó có tổng cộng 21 tỷ tham số, nhưng sử dụng 4 chuyên gia đang hoạt động cho mỗi token, trong tổng số 32 chuyên gia. Tính cả các thành phần được chia sẻ cộng với các chuyên gia đang hoạt động, mô hình này sử dụng ~3.6 tỷ tham số đang hoạt động cho mỗi token. Chạy mô hình này trên Mac M3 Ultra, có băng thông bộ nhớ khoảng 800 GB, chúng ta có thể ước tính tốc độ tạo là ~ 800 / (3.6 * 2) ở định dạng bfloat16, nơi mỗi tham số chiếm 2 byte. Điều này cho ra khoảng 111 token mỗi giây. Con số hiệu suất thực tế chúng ta nhận được là ~115 tok/s, rất gần với tính toán ước lượng.
Tốc độ siêu nhanh này xác nhận mô hình hoạt động gần như là một mô hình 3.6 tỷ tham số, nhưng nó có năng lực tương đương (hoặc chất lượng) như một mô hình 21 tỷ tham số.
(Lưu ý: Tốc độ sẽ còn nhanh hơn nữa nếu chúng ta sử dụng các kernel cho lượng tử hóa mxfp4 gốc mà mô hình sử dụng).
MoEs hấp dẫn vì những lý do sau:
-
Hiệu quả tính toán tốt hơn Với một ngân sách FLOP đào tạo cố định, MoEs thường vượt trội hơn các đối thủ dày đặc.
-
Một Trục Song Song Tự nhiên Các chuyên gia cung cấp một ranh giới cấu trúc trong đồ thị tính toán. Vì các token khác nhau sử dụng các chuyên gia khác nhau, chúng ta có thể song song hóa trên các chuyên gia (chúng ta sẽ thảo luận điều này sau trong phần Expert Parallelism).
-
Áp dụng trong Ngành Các bản phát hành MoE lớn gần đây của các mô hình mở trong vài tuần qua bao gồm Qwen 3.5, MiniMax M2, GLM-5, hoặc Kimi K2.5. Xu hướng này tăng tốc sau thành công của DeepSeek R1 vào tháng 1 năm 2025, dựa trên các hệ thống trước đó như DeepSeek V2. Một MoE sớm khác là Mixtral-8x7B, được phát hành vào tháng 12 năm 2023.
Transformers and MoEs
Hầu hết các công cụ trong hệ sinh thái, bao gồm tải mô hình, đặt thiết bị, lượng tử hóa và thực thi backend, ban đầu được thiết kế cho các mô hình dày đặc. MoEs thách thức những giả định này.
Việc biến MoEs thành công dân hạng nhất trong transformers có nghĩa là thiết kế lại các phần của quy trình tải, mô hình thực thi và các trừu tượng phân tán, chứ không chỉ thêm các lớp mô hình mới. Chúng ta sẽ tập trung vào cách thư viện transformers đã phát triển để hỗ trợ các kiến trúc thưa thớt trên:
- Tái cấu trúc tải trọng
- Backend Chuyên gia
- Song song hóa Chuyên gia
- Đào tạo MoE với transformers
Weight Loading Refactor
AutoModelForCausalLM.from_pretrained(“model_id”) tải xuống và tải trọng mô hình vào một mô hình PyTorch. Đối với các mô hình dày đặc, việc tải là tương đối đơn giản, trong đó mỗi tensor trong checkpoint ánh xạ một-đến-một với một tham số trong mô-đun thời gian chạy.
Đối với MoEs, nó phức tạp hơn. Trong hầu hết các checkpoint MoE, mỗi chuyên gia được tuần tự hóa độc lập. Nếu bạn xem bên trong chỉ mục checkpoint DeepSeek-V3, bạn sẽ thấy các khóa như:
bash model.layers.3.mlp.experts.0.gate_proj.weight … model.layers.3.mlp.experts.255.gate_proj.weight
Mỗi chuyên gia có bộ ma trận trọng số riêng, về cơ bản là 256 (tổng cộng từ 0 đến 255, lấy ví dụ DeepSeek-V3) các mạng chuyển tiếp nhỏ được lưu trữ cạnh nhau. Tuy nhiên, tại thời gian chạy, các kernel GPU được tối ưu hóa. Các kernel MoE hiện đại như GEMM nhóm và các triển khai MoE hợp nhất được thiết kế để xử lý tất cả các chuyên gia trong một thao tác duy nhất, chứ không phải bằng cách lặp qua chúng từng cái một.
Để làm điều đó một cách hiệu quả, chúng đòi hỏi các trọng số chuyên gia phải được đóng gói vào một tensor liên tục duy nhất.
Vì vậy, chúng ta có một sự không khớp:
- Checkpoint: 256 tensor riêng biệt
- Thời gian chạy: 1 tensor được đóng gói
Việc kết nối khoảng trống này một cách có hệ thống là điều mà tái cấu trúc tải trọng cho phép.
Với sự ra đời của một WeightConverter chung, mô hình tư duy đã chuyển từ:
thành:
Dynamic Weight Loading with WeightConverter
Trừu tượng trung tâm được giới thiệu bởi việc tái cấu trúc này là tải trọng động thông qua một WeightConverter.
WeightConverter cho phép chúng tôi định nghĩa:
python source key patterns → target key(s) + operations
Các thao tác nguyên thủy (chunk, concatenate, v.v.) có thể được kết hợp. Hai thao tác đặc biệt hữu ích cho MoEs:
-
MergeModulelist hợp nhất một danh sách các tensor thành một tensor duy nhất. Ví dụ, bạn có thể kết hợp MergeModulelist với Concatenate để xếp chồng các chuyên gia trong MoE và đóng gói chúng thành một tensor.
python WeightConverter( [“block_sparse_moe.experts..w1.weight”, “block_sparse_moe.experts..w3.weight”,], “mlp.experts.gate_up_proj”, operations=[ MergeModulelist(dim=0), Concatenate(dim=1), ], )
-
SplitModulelist tách một tensor trở lại thành một danh sách các tensor. Ví dụ, bạn có thể tách một chồng chuyên gia trở lại thành các chuyên gia riêng lẻ.
python WeightConverter( “mlp.experts.down_proj”, “block_sparse_moe.experts.*.w2.weight”, operations=[SplitModulelist(dim=0)], )
Lazy Materialization of Tensors
Việc tái cấu trúc không chỉ cải thiện những gì các chuyển đổi tồn tại, mà còn cách thức chúng được lên lịch.
Bộ tải quét các khóa checkpoint một lần, khớp chúng với các mẫu bộ chuyển đổi và nhóm các tensor cho mỗi bộ chuyển đổi. Một khi một khóa được xác định là cần thiết, nó sẽ được đăng ký dưới dạng một tương lai và được hiện thực hóa thông qua một nhóm luồng. Các hoạt động chuyển đổi chỉ chạy khi các phụ thuộc của chúng sẵn sàng. Ví dụ, MergeModulelist đợi cho đến khi tất cả các chuyên gia cho một lớp được tải.
Điều này tránh quét lại và giảm các đỉnh bộ nhớ.
Benchmark: Weight-Loading Pipeline Improvements
Để đánh giá các cải tiến được giới thiệu bởi quy trình tải trọng mới, chúng tôi đã so sánh hiệu suất giữa các phiên bản v4 và v5 của transformers. Trọng tâm là tốc độ tải các mô hình MoE lớn, đây thường là nút thắt cổ chai trong quá trình đào tạo và suy luận.
Chúng tôi đã so sánh hiệu suất v4 so với v5 sử dụng:
- Nhánh v4: https://github.com/ariG23498/transformers/tree/bench-v4
- Nhánh v5: https://github.com/ariG23498/transformers/tree/bench-v5
Ví dụ:
python from transformers import AutoModelForCausalLM
model_id = “Qwen/Qwen1.5-110B-Chat” model = AutoModelForCausalLM.from_pretrained(model_id)
Hai biến môi trường liên quan:
- HF_ENABLE_PARALLEL_LOADING: Kích hoạt tải shard song song thông qua các luồng.
- HF_DEACTIVATE_ASYNC_LOAD: Vô hiệu hóa quy trình bất đồng bộ mới (lối thoát v5).
Results
Model: Qwen/Qwen1.5-110B-Chat GPU: 1× A100 (80GB)
| Version | Strategy | Loading Mode | Time |
|---|---|---|---|
| v4.57.6 | device_map="auto" |
Threadpool | 66.24s |
| v4.57.6 | device_map="auto" |
Sequential | 67.29s |
| v4.57.6 | TP | — | OOM |
| v5 | device_map="auto" |
Async (default) | 20.71s |
| v5 | device_map="auto" |
Sync | 45.3s |
| v5 | TP | Async | 10.1s |
| v5 | TP | Sync | 19.28s |
Tốc độ tăng lên không chỉ đơn thuần là “nhiều luồng hơn”.
Đó là sự kết hợp của Định tuyến một lượt, Hiện thực hóa bất đồng bộ, và Lập lịch nhận biết chuyển đổi cùng nhau tránh hiện thực hóa không cần thiết và đỉnh bộ nhớ, đồng thời cho phép đóng gói chuyên gia và hợp nhất phép chiếu tại thời điểm tải.
Where Quantization Fits In
Với việc tái cấu trúc này, chúng ta có thể tạo cấu trúc mô hình thời gian chạy trước và sau đó chuyển đổi trọng số vào cấu trúc đó. Chúng ta có thể tùy chọn đính kèm lượng tử hóa trong quy trình chuyển đổi, làm cho lượng tử hóa trở thành một phần của chính quy trình tải trọng. Điều này rất quan trọng vì việc lượng tử hóa “theo từng chuyên gia” chỉ có ý nghĩa khi các chuyên gia tồn tại trong bố cục đóng gói có thể dự đoán được.
Quy trình đầu cuối này trước đây không khả thi và bây giờ nó được cung cấp cho người dùng dưới dạng một API hiển thị.
Expert Backend
Khi các chuyên gia đã được đóng gói vào một tensor thời gian chạy duy nhất, một câu hỏi khác nảy sinh:
Trong mô hình Mixture of Experts, mỗi token được định tuyến đến các chuyên gia khác nhau. Điều này có nghĩa là thời gian chạy phải điều phối các token đến các trọng số chuyên gia đã chọn của chúng, thực thi các phép chiếu một cách hiệu quả, áp dụng các trọng số định tuyến và sau đó thu thập và sắp xếp lại kết quả.
Đây là những gì Hệ thống Backend Chuyên gia (được giới thiệu trong PR #42697) giải quyết. Backend Chuyên gia giới thiệu một kiến trúc thực thi có thể cắm được tách biệt tính toán chuyên gia khỏi việc triển khai mô hình. Thay vì mã hóa cứng một chiến lược điều phối duy nhất bên trong mỗi mô hình MoE, hệ thống cho phép các lớp chuyên gia chọn backend một cách động tại thời gian chạy.
Điều này được thực hiện thông qua một mẫu decorator:
python @use_experts_implementation
Decorator này gói các lớp chuyên gia và tự động điều phối tính toán đến backend đã chọn.
Ba backend hiện được cung cấp:
- eager: lặp qua các chuyên gia đã chọn và áp dụng các phép chiếu cho mỗi chuyên gia. Backend này được sử dụng để tham chiếu độ chính xác và gỡ lỗi.
- batched_mm: sử dụng API torch.bmm. Backend này nhân đôi các trọng số chuyên gia đã chọn cho mỗi token và thực hiện một phép GEMM theo lô duy nhất. Backend này rất phù hợp cho các tải trọng nặng GPU, lô nhỏ, nơi có bộ nhớ.
- grouped_mm: sử dụng API torch._grouped_mm. Ở đây, chúng tôi sắp xếp các token theo ID chuyên gia, nhóm chúng và sau đó thực hiện một phép GEMM theo nhóm duy nhất. Backend này tỏa sáng với các lô lớn hoặc các thiết lập bị giới hạn bộ nhớ.
Expert Parallelism
Mixture of Experts (MoE) có thể có hàng trăm tỷ tham số (nhiều hơn nhiều so với những gì có thể nằm trên một GPU duy nhất). Song song hóa chuyên gia (EP) giải quyết vấn đề này bằng cách phân phối các chuyên gia trên nhiều thiết bị. Mỗi thiết bị chỉ tải tập hợp con các chuyên gia được chỉ định, tính toán cho các chuyên gia đó và sau đó tham gia vào việc tổng hợp kết quả. Phương pháp này mở rộng quy mô mô hình đế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 chuyên gia.
Song song hóa chuyên gia được kích hoạt thông qua enable_expert_parallel:
python import torch from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.distributed.configuration_utils import DistributedConfig
distributed_config = DistributedConfig(enable_expert_parallel=True)
model = AutoModelForCausalLM.from_pretrained( “openai/gpt-oss-120b”, dtype=“auto”, distributed_config=distributed_config, )
Khởi chạy với:
bash torchrun –nproc-per-node N script.py
Ở đây N chia đều cho tổng số chuyên gia, và có thể khớp với số lượng GPU trên nút của bạn.
Khi enable_expert_parallel=True, mô hình chuyển từ kế hoạch song song hóa tensor (TP) tiêu chuẩn sang kế hoạch song song hóa chuyên gia (EP) với các chiến lược phân chia chuyên dụng.
Các thành phần cốt lõi của EP nằm trong:
- GroupedGemmParallel: Điều này tách các trọng số chuyên gia dọc theo chiều chuyên gia (dim=0). Ở đây mỗi thiết bị chỉ tải num_experts / num_devices.
- RouterParallel: Điều này ánh xạ lại các chỉ số chuyên gia toàn cục sang các chỉ số cục bộ, loại bỏ các chuyên gia không được gán cho thứ hạng hiện tại, đảm bảo mỗi thiết bị chỉ tính toán với các chuyên gia cục bộ của nó và sử dụng tất cả các giảm để kết hợp các kết quả đầu ra một phần trên các thiết bị.
Training MoEs with Transformers
MoEs rất tuyệt vời để mở rộng quy mô suy luận, nhưng việc đào tạo chúng phức tạp hơn đáng kể.
MoEs có số lượng tham số khổng lồ, giao tiếp chuyên gia phân tán rất phức tạp, có sự bất ổn định trong định tuyến cần được xử lý. Để giải quyết vấn đề này, chúng tôi đã hợp tác với Unsloth để cho phép đào tạo Mixture-of-Experts nhanh hơn đáng kể:
- ~12 lần đào tạo MoE nhanh hơn
- Giảm >35% VRAM
- Ngữ cảnh dài hơn ~6 lần
- Tăng tốc tổng thể gấp 12–30 lần so với v4
Chúng tôi tận dụng trừu tượng Backend Chuyên gia, tiêu chuẩn hóa trên API torch._grouped_mm của PyTorch và sử dụng các kernel Triton grouped-GEMM + LoRA tùy chỉnh. Unsloth xây dựng dựa trên các tối ưu hóa của Transformers (và TRL) để đẩy hiệu suất xa hơn.
Conclusion
Khi các kiến trúc thưa thớt tiếp tục phát triển, chúng tôi muốn thư viện transformers phát triển cùng chúng. Nếu bạn đang xây dựng với MoEs hoặc thử nghiệm các ý tưởng thưa thớt mới, chúng tôi rất muốn nghe ý kiến của bạn. Hãy cho chúng tôi biết bạn muốn thấy những trừu tượng, kernel hoặc quy trình làm việc nào tiếp theo trong transformers.
More Articles from our Blog
- Continuous batching from first principles
- Tricks from OpenAI gpt-oss YOU 🫵 can use with transformers