Tăng tốc huấn luyện mô hình MoE không loại bỏ (dropless) trong JAX với NVIDIA Transformer Engine

Mixture of experts (MoE) đã trở thành một trong những xu hướng kiến trúc mang tính định hình trong việc huấn luyện các mô hình AI quy mô lớn. DeepSeek, Qwen và Mixtral là những ví dụ về các mô hình MoE có hiệu suất ngang bằng hoặc vượt trội so với các mô hình dense tương đương nhưng chỉ tốn một phần nhỏ tài nguyên tính toán khi huấn luyện.

Các mô hình MoE cung cấp khả năng huấn luyện hiệu quả thông qua tính toán có điều kiện (conditional computation). Thay vì sử dụng một mạng feed-forward network (FFN) dày đặc dùng chung cho tất cả các token, MoE thay thế nó bằng nhiều mạng chuyên gia (expert networks) nhỏ hơn và một bộ định tuyến (router) được học để quyết định kích hoạt Top-K chuyên gia nào.

Tuy nhiên, việc tối ưu hóa huấn luyện MoE ở quy mô lớn là một thách thức. Trong quá trình huấn luyện DeepSeek-V3 trên NVIDIA GB200, một baseline chưa tối ưu chỉ đạt 103 TFLOPS/GPU, với giao tiếp liên GPU chiếm tới 84% thời gian kernel tích lũy. Với thư viện JAX Python và các tối ưu hóa kernel nhắm mục tiêu của NVIDIA Transformer Engine, con số đó đã tăng lên 1.068 TFLOPS/GPU, cải thiện gấp 10,4 lần. Bài viết này thảo luận về cách Transformer Engine – một thư viện để tăng tốc các mô hình Transformer trên GPU NVIDIA – kết hợp với JAX mang lại sự cải thiện hiệu suất đáng kể trong các hoạt động của mô hình MoE.

Những thách thức trong huấn luyện MoE là gì?

Huấn luyện MoE ở quy mô sản xuất tạo ra các nút thắt cổ chai không tồn tại ở các mô hình dense: định tuyến token (token routing), điều phối và tập hợp chuyên gia (expert dispatch and gather), giao tiếp all-to-all và các phép toán expert GEMM không đồng nhất (ragged expert GEMMs).

Vấn đề trở nên phức tạp hơn vì bộ định tuyến là thành phần được học. Trong suốt quá trình huấn luyện, phân phối có thể trở nên lệch lạc đáng kể khi bộ định tuyến phát triển sự ưu tiên cho một số chuyên gia nhất định. Không có hai batch nào tạo ra tải chuyên gia giống nhau, và trong một batch, một chuyên gia có thể nhận được nhiều token hơn hẳn các chuyên gia khác. Mỗi chuyên gia nhận được số lượng token khác nhau, vì vậy không có cấu trúc GEMM hình chữ nhật sạch để batch và điều phối. Điều này dẫn đến các tensor không đồng nhất (ragged tensors).

Trong MoE, các token được định tuyến linh hoạt đến các chuyên gia khác nhau. Điều này có nghĩa là số lượng token được gán cho mỗi chuyên gia thay đổi một cách khó dự đoán, dẫn đến các tensor không đồng nhất (Hình 1). Đây là một thách thức vì hầu hết các thư viện đều được tối ưu hóa cao cho các phép toán tensor yêu cầu cấu trúc dữ liệu đồng nhất, hình chữ nhật.

Với expert parallelism (EP), các token phải được điều phối và các đầu ra phải được kết hợp, khôi phục về thứ tự token ban đầu. Nếu đường dẫn điều phối và kết hợp không được tối ưu hóa, giao tiếp sẽ chiếm ưu thế và GPU bị sử dụng dưới mức công suất. Một cơ chế all-to-all được tối ưu hóa kém sẽ buộc GPU phải tạm dừng và chờ dữ liệu trước khi thực hiện bất kỳ công việc hữu ích nào.

Giải quyết vấn đề này đòi hỏi các kernel chuyên dụng có khả năng xử lý nguyên bản các bố cục không đồng nhất. Đây chính xác là vấn đề mà các tối ưu hóa MoE của Transformer Engine được thiết kế để giải quyết.

Dropless MoE khác gì với capacity-based MoE?

Dropless MoE và capacity-based MoE là hai cách tiếp cận khác nhau để xử lý việc định tuyến token đến các chuyên gia.

Trong dropless MoE, mọi token đều được xử lý bởi chuyên gia đã chọn bất kể tải trọng không đồng đều ra sao. Điều này hấp dẫn đối với chất lượng mô hình nhưng lại đòi hỏi khắt khe về hệ thống. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts đã giải quyết vấn đề này bằng cách tái cấu trúc phép tính chuyên gia thành phép nhân ma trận thưa khối (block-sparse matrix multiplication), cho phép mỗi chuyên gia hoạt động trên số lượng token khác nhau mà không cần loại bỏ (dropping) hoặc đệm (padding). Điều này đòi hỏi các kernel GPU block-sparse mới, grouped GEMM được tối ưu hóa, cùng các primitive điều phối và kết hợp được thiết kế riêng cho số lượng token biến đổi.

Ngược lại, các khung huấn luyện MoE dựa trên dung lượng (capacity-based) tiêu chuẩn tránh được sự phức tạp của định tuyến động bằng cách hạn chế nó. Mỗi chuyên gia được gán một ngân sách token cố định, và bất kỳ phần dư thừa nào đều bị cắt bỏ hoặc đệm để vừa vặn. Điều này giữ cho tính toán ổn định và thân thiện với phần cứng, nhưng nó buộc phải đánh đổi trực tiếp giữa chất lượng mô hình và hiệu suất: loại bỏ các token dư thừa thì mô hình huấn luyện trên dữ liệu không đầy đủ, hoặc đệm để tránh loại bỏ thì phải trả giá bằng việc lãng phí tài nguyên tính toán và bộ nhớ.

Những tối ưu hóa chuyên dụng nào là cần thiết cho dropless MoE?

Cam kết sử dụng dropless MoE đồng nghĩa với việc stack huấn luyện không còn có thể dựa vào các hình dạng chuyên gia cố định. Mọi kernel liên quan đến tính toán chuyên gia đều phải xử lý số lượng token biến đổi một cách hiệu quả. Ngoài ra, điều này có nghĩa là số lượng token của mỗi chuyên gia là biến đổi và phụ thuộc vào dữ liệu, vì vậy các kernel không chỉ phải chấp nhận các hình dạng động mà còn phải hoạt động khi các hình dạng đó không thể truy cập được trên CPU để kích hoạt CUDA graphs và tránh việc biên dịch lại.

Transformer Engine cung cấp các khối xây dựng sau giúp cách tiếp cận này trở nên thực tế trong JAX:

Hình 3 cho thấy một lớp MoE song song chuyên gia (expert-parallel) trên hai GPU. Bộ định tuyến gán mỗi token cho một chuyên gia, quá trình dispatch di chuyển các token đến GPU của chuyên gia đó. Grouped MLP chạy hai grouped GEMM trên các nhóm có độ dài biến đổi đó, và combine đảo ngược quá trình trao đổi để khôi phục thứ tự token ban đầu.

Tối ưu hóa 1: Grouped GEMM

Trong một FFN dày đặc, mọi token đều đi qua cùng một ma trận trọng số. Trong MoE, bộ định tuyến phân phối token không đồng đều nên mỗi chuyên gia nhận được số lượng token khác nhau mỗi bước, phá vỡ hình dạng GEMM thông thường mà các kernel điển hình được tối ưu hóa.

Các phương pháp trước đây bao gồm một vòng lặp các kernel GEMM và batched GEMM. Vòng lặp yêu cầu các bản sao Device-to-Host của số lượng token. Điều này nằm trên đường dẫn quan trọng (critical path), gây ra độ trễ của việc truyền Device-to-Host và phá vỡ CUDA graphs. Batched GEMM tính toán dung lượng token trường hợp xấu nhất ngay cả khi sử dụng ít token hơn vì chúng được đệm để buộc tính toán chuyên gia cố định, dẫn đến tính toán dư thừa.

Một grouped GEMM giải quyết vấn đề này bằng cách xử lý tất cả các phép nhân ma trận chuyên gia trong một lần gọi kernel duy nhất, mỗi phép toán với số lượng token thực tế của nó. Nó chỉ tính toán các vùng có token hợp lệ và kết quả là hiệu suất cao hơn.

Transformer Engine grouped_gemm / ragged_dot hỗ trợ điều này với cuBLAS và cuBLASLt, ánh xạ trực tiếp vào các thư viện GEMM hiệu suất tốt nhất của NVIDIA để mang lại khả năng sử dụng Tensor Core tối đa ngay cả với các hình dạng chuyên gia không đều. Trên GPU NVIDIA Blackwell, đường dẫn này cũng mở ra khả năng mở rộng khối MXFP8 cho các phép nhân ma trận chuyên gia sử dụng các kernel lượng tử hóa nhóm của Transformer Engine.

Tối ưu hóa 2: Expert parallelism để tích hợp Dispatch và Combine

Sau khi các kernel bộ định tuyến hợp nhất gán mỗi token cho các chuyên gia của nó, mô hình phải di chuyển vật lý các token đó đến đúng thiết bị, xử lý chúng và đưa kết quả trở lại.

Quá trình này chia thành hai giai đoạn riêng biệt: Dispatch và Combine.

Trong một triển khai ngây thơ, các giai đoạn này chạy như một chuỗi nối tiếp các hoạt động riêng biệt, với việc GPU bị tạm dừng giữa các bước, dữ liệu bị đọc và ghi vào bộ nhớ nhiều lần, và giao tiếp hầu như không hoạt động trong khi tính toán chạy và ngược lại.

Triển khai EP của Transformer Engine tích hợp các giai đoạn Dispatch và Combine vào một đường dẫn kernel được hợp nhất chặt chẽ. Sự tích hợp này được hỗ trợ bởi NCCL EP, một backend giao tiếp được tinh chỉnh đặc biệt cho các mô hình lưu lượng không đều, mất cân bằng mà định tuyến song song chuyên gia tạo ra.

NCCL EP cũng sử dụng cơ chế khử trùng lặp token (token deduplication): khi một token được điều phối đến nhiều chuyên gia trên cùng một rank hoặc đến nhiều rank trên một node IB từ xa, nó chỉ truyền qua mạng một lần và được sao chép trên node nhận, giúp tiết kiệm băng thông mạng. EP là đối trọng của grouped GEMM: grouped GEMM xử lý những gì xảy ra bên trong mỗi chuyên gia; EP xử lý mọi thứ xung quanh nó.

Các tối ưu hóa bổ sung

Các tối ưu hóa bổ sung bao gồm JAX host offloading và XLA multistreaming collectives.

JAX host offloading

Các kích hoạt trung gian (intermediate activations) không cần phải được lưu trên thiết bị trong toàn bộ quá trình forward pass. JAX cung cấp các API tái tạo (rematerialization) để offload các kích hoạt sang bộ nhớ host. Để tiết kiệm bộ nhớ trong quá trình huấn luyện DSv3, hãy offload kết quả của query và value projection sang host. Để tìm hiểu thêm, hãy xem Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading.

XLA multistreaming collectives

Trong khi EP được điều khiển bởi Transformer Engine NCCL EP, FSDP được tối ưu hóa được xử lý nguyên bản trong XLA. Theo mặc định, XLA chạy giao tiếp trên một luồng duy nhất, vì vậy các tập thể (collectives) có thể thực thi song song lại bị tuần tự hóa và một số cuối cùng bị lộ trên đường dẫn quan trọng. Multi-stream collectives cho phép trình biên dịch lập lịch các tập thể độc lập đồng thời trên các luồng CUDA riêng biệt, chồng lấp các lần truyền InfiniBand liên node với giao tiếp NVIDIA NVLink nội node để tận dụng cả hai cấu trúc cùng một lúc thay vì chờ đợi trên một luồng tuần tự hóa.

Latency Hiding Scheduler (LHS) quyết định tập thể nào an toàn để chồng lấp bằng cách phân tích các nhóm bản sao của chúng và kiểm tra rủi ro bế tắc (deadlock), vì vậy mức tăng băng thông bộ nhớ là tự động và không yêu cầu chú thích thủ công. Điều này làm giảm đáng kể tỷ lệ phần trăm các tập thể bị lộ trong quá trình huấn luyện DSv3.

Tác động hiệu suất huấn luyện của MoE trong JAX với Transformer Engine là gì?