Giải mã nhân EP hiệu suất cao: Cốt lõi của việc chạy mô hình MoE quy mô lớn

Các mô hình ngôn ngữ lớn (LLM) ngày nay có kích thước khổng lồ, đòi hỏi hàng loạt GPU để vận hành. Trong lý tưởng, việc suy luận LLM sẽ hoàn toàn song song hóa, cho phép chúng ta tính toán các tác vụ độc lập trên các GPU khác nhau. Tuy nhiên, thực tế phức tạp hơn nhiều: để sử dụng nhiều GPU cho LLM, chúng ta cần thiết lập giao tiếp giữa chúng.

Có nhiều cách để kết nối các GPU với nhau như: Song song hóa Tensor (Tensor Parallelism), Song song hóa Pipeline (Pipeline Parallelism), hay Song song hóa Ngữ cảnh (Context Parallelism). Nhưng đối với các mô hình Mixture of Experts (MoE), khi cần phục vụ ở quy mô lớn, "Song song hóa Chuyên gia rộng" (wide Expert Parallelism - wideEP) đóng vai trò quan trọng nhất.

Bài viết này sẽ phân tích giải phẫu của một nhân EP hiệu suất cao theo phong cách DeepEP, tập trung vào hai chiến lược chính: tối ưu hóa độ trục xuất (throughput) và tối ưu hóa độ trễ (latency).

Thách thức của việc định tuyến Token

Trong các mô hình MoE, các lớp chuyên gia (MoE layers) yêu cầu các token phải di chuyển đến đúng nơi các chuyên gia đó đang cư ngụ. Khác với các hình thức song song hóa khác có mô hình giao tiếp cố định, Expert Parallelism (EP) rất linh hoạt. Việc quyết định token nào cần đến GPU nào do bộ định tuyến (router) quyết định dựa trên dữ liệu tại thời điểm chạy (runtime).

Ví dụ, với một cụm 8 GPU chia thành 2 nút, mỗi token có thể được định tuyến đến 2 chuyên gia khác nhau. Nhiệm vụ của nhân giao tiếp EP là đảm bảo các token này gặp gỡ chuyên gia của mình bất kể họ nằm ở đâu trong mạng, thực hiện tính toán GEMM (nhân ma trận) và sau đó đưa kết quả về lại vị trí ban đầu.

Tối ưu hóa độ trục xuất: Hỏi trước, gửi sau

Chiến lược này phù hợp nhất cho giai đoạn prefill (giai đoạn nạp đầu vào), nơi có các lô (batch) lớn và nhiều tác vụ tính toán khác để che giấu thời gian giao tiếp.

Phân phối (Dispatch)

Mục tiêu là cung cấp dữ liệu cho một "grouped GEMM" – một phép nhân ma trận chạy riêng cho từng chuyên gia trên số lượng token tương ứng. Thách thức nằm ở việc chúng ta không biết trước kích thước của bộ đệm cần thiết, vì số lượng token mỗi chuyên gia nhận được thay đổi liên tục.

Để tối ưu hóa bộ nhớ HBM (High Bandwidth Memory), chúng ta không thể cấp phát bộ đệm cho trường hợp xấu nhất (worst-case) vì sẽ lãng phí tài nguyên. Thay vào đó, phương pháp độ trục xuất cao sử dụng một bước phối hợp (coordination pass).

1. Trao đổi số lượng: Mỗi rank (GPU) trao đổi số lượng token họ định gửi cho các peer khác. Qua đó, mỗi rank biết chính xác họ sẽ nhận bao nhiêu token.
2. Tính toán địa chỉ: Sử dụng tổng tiền tố (prefix sum), hệ thống xác định vị trí chính xác của từng khối dữ liệu trong bộ đệm nén (compact buffer).
3. Gửi dữ liệu: Người gửi đẩy token vào một hàng đợi cố định trên đích, và người nhận kéo dữ liệu từ hàng đợi đó vào vị trí đã tính toán.

Cuối cùng, một bước hoán vị cục bộ (local permute) sẽ sắp xếp dữ liệu từ thứ tự "theo nguồn" sang thứ tự "theo chuyên gia" để GEMM có thể xử lý.

Kết hợp (Combine)

Đây là quá trình ngược lại của Dispatch. Sau khi GEMM hoàn tất, đầu ra được sắp xếp lại theo thứ tự nguồn ban đầu và gửi trở lại rank gốc. Tại đây, các đầu ra riêng lẻ từ các chuyên gia khác nhau được cộng lại (weighted sum) để tạo ra vector đầu ra cuối cùng cho token đó.

Tối ưu hóa độ trễ: Gửi ngay mà không cần hỏi

Trong giai đoạn decode (giải mã), mỗi rank chỉ giữ một số ít token (thường là 1 token cho mỗi chuỗi). Bước phối hợp trong phương pháp trên trở nên quá tốn kém vì nó là một vòng lặp mạng đầy đủ. Chúng ta cần loại bỏ nó.

Chiến lược bộ đệm cố định

Thay vì hỏi trước để cấp phát bộ đệm nén, chúng ta dự trù trước một vùng nhớ cố định cho mỗi cặp (nguồn, chuyên gia). Điều này cho phép người gửi tính toán địa chỉ đích ngay lập tức bằng một công thức tĩnh mà không cần chờ đợi thông tin từ các rank khác.

Địa chỉ ghi bây giờ là một hàm số đơn giản, không phụ thuộc vào hành động của người gửi khác. Dữ liệu được gửi đi ngay lập tức.

Đánh đổi và tối ưu hóa

• Lãng phí bộ nhớ: Vì vùng nhớ được cấp phát cho trường hợp xấu nhất, bộ đệm sẽ phần lớn trống. Đây là cái giá phải trả để loại bỏ độ trễ của vòng lặp mạng.
• Tín hiệu hoàn thành: Để người nhận biết dữ liệu đã sẵn sàng mà không cần bước phối hợp, người gửi ghi thêm một giá trị "số lượng token" vào một vị trí cố định. Việc xuất hiện giá trị này cũng là tín hiệu báo rằng dữ liệu đã được ghi xong.
• Lượng tử hóa: Để giảm lượng byte truyền đi (giảm độ trễ), DeepEP sử dụng định dạng FP8 cho dữ liệu truyền qua mạng trong giai đoạn Dispatch, trong khi giai đoạn Combine vẫn giữ độ chính xác BF16 cho các phép cộng.

Kết luận

Nói tóm lại, tối ưu hóa độ trục xuất chấp nhận đánh đổi một vòng lặp mạng để tiết kiệm bộ nhớ HBM, trong khi tối ưu hóa độ trễ chấp nhận lãng phí bộ nhớ (thông qua các bộ đệm cố định trống) để loại bỏ vòng lặp đó.

Hiện nay, thư viện DeepEP đã chuyển sang sử dụng API giao tiếp phía thiết bị mới của NCCL. Các dự án như UCCL cũng đang tái hiện các nguyên thủy này cho GPU AMD và các loại NIC RDMA khác. Trong tương lai, chúng ta có thể mong đợi việc kết hợp (fuse) các nguyên thủy giao tiếp này trực tiếp vào nhân tính toán để đạt được sự chồng chéo tinh grained hơn nữa.