Giải mã nhân EP hiệu suất cao: Cốt lõi tối ưu cho LLM 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ồ. Chính vì thế, chúng ta cần hàng loạt GPU để chạy chúng. Sẽ rất tuyệt nếu việc suy luận (inference) của LLM có thể "song song một cách dễ dàng" (embarrassingly parallel), để chúng ta có thể 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 bắt buộc phải khiến các GPU đó giao tiếp với nhau.

Có nhiều cách để kết nối các GPU lại với nhau: Song song Tensor (Tensor Parallelism), Song song Pipeline (Pipeline Parallelism), Song song Ngữ cảnh (Context Parallelism), Song song Chuyên gia (Expert Parallelism), v.v. Mỗi phương pháp đều có vai trò riêng. Nhưng đối với các mô hình MoE (Mixture of Experts), ở các lớp MoE, khi muốn phục vụ ở quy mô lớn, "Wide Expert Parallelism" (wideEP) chính là vua.

Các loại song song khác đều yêu cầu giao tiếp giữa các GPU, nhưng mẫu hình giao tiếp của chúng được cố định bởi kiến trúc: ai gửi, ai nhận và gửi bao nhiêu đều được xác định trước khi forward pass bắt đầu. Tuy nhiên, Expert Parallelism thì khác. Token nào cần đến GPU nào được quyết định bởi bộ định tuyến (router), dựa trên dữ liệu, tại thời điểm chạy (runtime), và thay đổi ở mỗi lớp MoE.

Bài viết này sẽ đi sâu vào giải phẫu học của một nhân giao tiếp EP theo phong cách DeepEP: trước hết là hình dạng tối ưu hóa độ trục xuất (throughput), sau đó là hình dạng tối ưu hóa độ trễ (latency).

Nhiệm vụ của chúng ta

Hãy cụ thể hóa bài toán. Chúng ta có 8 GPU, chia cắt trên 2 node, kết nối qua RDMA, và mỗi rank song song dữ liệu sở hữu một GPU. Attention chạy trên mỗi GPU qua một lô token $B_i$, trong đó $B_i$ có thể thay đổi giữa các GPU. Chúng ta đang thực hiện expert parallel với $E=16$ chuyên gia, hai chuyên gia trên mỗi GPU, trong đó $K=2$ chuyên gia được định tuyến cho mỗi token.

Mục tiêu của các nhân song song chuyên gia là đưa các activation đến nơi chúng cần đến, chạy các phép nhân ma trận chuyên gia (expert GEMMs) khi đến đó, và sau đó đưa chúng về lại vị trí ban đầu.

Trong giao tiếp, chúng ta thường quan tâm đến hai yếu tố: độ trục xuất (throughput) hoặc độ trễ (latency). Sự phân chia này tương ứng với hai giai đoạn của suy luận: prefill mang lại các lô lớn, bị giới hạn bởi tính toán, có nhiều công việc khác để che giấu thời gian giao tiếp; trong khi đó, ở giai đoạn decode, hầu như không có việc gì khác để làm, nên việc truyền dữ liệu chính là thứ chúng ta phải chờ đợi.

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

Dispatch (Phân phối)

Mục đích của dispatch là cung cấp dữ liệu cho một GEMM nhóm (grouped GEMM). Sau khi định tuyến, mỗi chuyên gia phải chạy một phép nhân ma trận trên chính xác các token được gán cho nó. Trước khi định tuyến, các token này bị phân tán trên mọi rank trong cụm. Vì vậy, trên mỗi rank, dispatch phải thu thập các token được định hướng đến các chuyên gia địa phương vào một bộ đệm duy nhất để GEMM nhóm có thể tiêu thụ một lần.

Khó khăn nằm ở chỗ chúng ta không biết hình dạng của bộ đệm đó trước. Token nào đi đến chuyên gia nào được quyết định bởi router tại thời điểm chạy, và sự phân bố này không đồng đều và thay đổi ở mỗi bước. Một chuyên gia có thể thu hút hàng nghìn token lúc này và không có token nào lúc khác.

Có hai cách để xử lý việc không biết kích thước:

1. Dự phòng đủ chỗ cho trường hợp xấu nhất bằng cách cấp phát một hình chữ nhật cố định.
2. Tìm hiểu số lượng thực tế trước rồi cấp phát chính xác những gì chúng ta cần.

Hình chữ nhật cố định đơn giản hơn nhưng phải được định kích thước cho trường hợp xấu nhất. Điều này lãng phí bộ nhớ HBM. Trong khi đó, chi phí để biết số lượng thực tế là chấp nhận được ở giai đoạn prefill, vì các lô lớn và GEMM bị giới hạn bởi tính toán cho phép ẩn đi thời gian giao tiếp thêm. Do đó, đường dẫn độ trục xuất cấp phát một bộ đệm không đều (ragged buffer), được định kích thước theo số lượng token thực tế sẽ nhận.

Nếu muốn cấp phát chính xác những gì cần thiết, chúng ta phải biết số lượng trước khi bất kỳ activation nào di chuyển. Chúng ta có thể thực hiện một bước phối hợp (coordination pass). Mọi rank đều biết từ định tuyến của chính mình rằng họ đang gửi bao nhiêu token đến từng peer. Nếu mọi người trao đổi các con số này, mỗi rank có thể cộng tổng xem họ sắp nhận bao nhiêu.

Bước phối hợp này rẻ về byte, chỉ là một vài số nguyên cho mỗi peer thay vì megabyte trạng thái ẩn. Khi một rank biết bao nhiêu token đang đến từ mỗi nguồn, bố cục bộ đệm an toàn để ghi sẽ tự nhiên xuất hiện dưới dạng tổng tiền tố (prefix sum).

Với bố cục đã cố định, chúng ta có thể thực sự gửi các activation. Người gửi không bao giờ ghi trực tiếp vào bộ đệm cuối cùng. Thay vào đó, người gửi truyền các token của mình vào một hàng đợi kích thước nhỏ cố định ở đích, được khắc ra từ bộ nhớ đã đăng ký trước. Người nhận, người sở hữu bộ đệm nhỏ gọn, sẽ làm rỗng hàng đợi đó và sao chép từng token vào vị trí mà tổng tiền tố đã chỉ định.

Combine (Tổng hợp)

Mục đích của combine là hoàn tác nhân dispatch và cộng các đóng góp cho mỗi token.

GEMM để lại các đầu ra được nhóm theo chuyên gia, vì vậy điều đầu tiên chúng ta phải làm là hoàn tác sự hoán vị mà chúng ta đã thực hiện trên đường đi. Hoán vị ngược đưa các đầu ra trở lại bố cục theo nguồn rank mà dispatch đã phân phối token.

Từ đó, quá trình truyền tải chạy theo chiều ngược lại. Rank lưu trữ chuyên gia giờ đây là người gửi, truyền các đầu ra của nó trở lại qua các hàng đợi theo từng peer của rank gốc về các vị trí mà các token đã đến.

Chúng ta không cần thực hiện bước phối hợp, vì combine được cung cấp cùng thông tin định tuyến mà dispatch tạo ra, vì vậy nó đã biết mọi thứ cần trả về đâu.

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

Lý do để tối ưu hóa nhân cho độ trễ là chế độ decode. Mỗi rank chỉ giữ một vài token, thường là một token cho mỗi chuỗi trong lô. Bước phối hợp rẻ về byte, nhưng đó là một vòng lặp mạng đầy đủ với các rào cản (barriers), và nó phải hoàn thành trước khi bất kỳ activation nào di chuyển. Ở decode, có rất ít việc để chồng chéo lên nó, và nó trở thành một phần lớn của lớp.

Vì vậy, chúng ta muốn tìm cách bỏ qua nó. Nếu chúng ta sẵn sàng từ bỏ tính nhỏ gọn, chúng ta có thể sắp xếp trước không gian cho mỗi peer rank để ghi vào. Thay vì một bộ đệm đóng gói, chúng ta dự trữ trước một vùng riêng tư cố định cho mỗi cặp (rank nguồn, chuyên gia). Đây là hình chữ nhật cố định từ nhánh dispatch, với một sự tinh chỉnh: phần đệm là theo từng (rank nguồn, chuyên gia) thay vì theo từng chuyên gia, để không có hai người gửi nào ghi vào cùng một vùng.

Địa chỉ mà người gửi ghi vào giờ đây là một công thức cố định. Tổng tiền tố động trên số lượng thực tế trở thành sải bước tĩnh nhân với mức tối đa cố định, và điều đầu tiên xảy ra trong lớp là việc gửi dữ liệu.

Vì việc truyền tải giờ đây là thứ chúng ta chờ đợi, số byte chính là độ trễ, và dispatch độ trễ thấp của DeepEP giảm chúng bằng cách lượng tử hóa tải trọng (payload) thành FP8 trên đường dây theo mặc định. Đường dẫn ngược lại giữ nguyên BF16: các tổng của combine là nơi độ chính xác quan trọng.

Điểm bắt bẻ là mỗi vùng riêng tư phải đủ lớn cho bất kỳ điều gì mà nguồn của nó có thể gửi, được định kích thước cho trường hợp xấu nhất thay vì số lượng thực tế. Nếu không bị giới hạn, trường hợp xấu nhất đó là khổng lồ. Vì vậy, chúng ta cần giới hạn số lượng token một rank có thể dispatch trong một lệnh gọi xuống một kích thước chunk cố định.

Combine hoạt động theo cách tương tự. Nó không bao giờ cần bước phối hợp, nhưng đường dẫn độ trục xuất vẫn dàn dựng các lượt trả về của nó qua các hàng đợi; ở đây ngay cả những hàng đợi đó cũng biến mất. Dispatch đã giao từng token được gắn thẻ nơi nó đến từ đâu, vì vậy máy chủ của chuyên gia có thể tính toán trực tiếp địa chỉ trả về: một vị trí riêng tư trên rank gốc của token, được lập chỉ mục bởi vị trí của token ở đó và chuyên gia K nào là chuyên gia này.

Tóm lại: độ trễ thấp đảo ngược sự đánh đổi của độ trục xuất. Độ trục xuất tốn một vòng lặp mạng để giữ bộ nhớ chặt chẽ; độ trễ thấp tốn bộ nhớ, dưới dạng bộ đệm trường hợp xấu nhất hầu hết trống rỗng, để loại bỏ vòng lặp mạng.

Kết luận

Đó là câu chuyện. Nếu bạn mở mã nguồn của DeepEP, bạn sẽ nhận ra hình dạng này, mặc dù gần đây bạn sẽ tìm thấy các nhân này dưới thư mục legacy/: bản viết lại V2 gần đây đã xây dựng lại thư viện trên top của API giao tiếp phía thiết bị mới của NCCL.

DeepEP được xây dựng cho ngăn xếp của NVIDIA: GPU Hopper và Blackwell, NVLink trong node, và RDMA lớp InfiniBand giữa các node. Dự án UCCL triển khai lại các nguyên thủy tương tự cho nhiều nền tảng hơn: GPU AMD cũng như NVIDIA, và bất kỳ NIC RDMA nào, AWS EFA, Broadcom, với hiệu suất tương đương.

Có một ngăn xếp tối ưu hóa đang phát triển ở trên. Cân bằng tải chuyên gia (EPLB) tính toán các kế hoạch sao chép và đặt vị trí cho các chuyên gia nóng, mà các hệ thống phục vụ áp dụng định kỳ. vLLM's elastic EP mở rộng và thu hẹp việc triển khai. Và các thống kê định tuyến có thể quan sát được chính tại lớp này.

Công việc đang được tiến hành để hợp nhất các nguyên thủy giao tiếp loại này vào chính các nhân tính toán, để chúng ta có thể thực hiện chồng chéo hạt mịn và đường ống tốt hơn: một SM có thể nhận dữ liệu trong khi các ô GEMM bắt đầu bắn ra từ nó.

Tuy nhiên, ranh giới đó thay đổi như thế nào, công việc vẫn là công việc mà chúng ta đã bắt đầu: các token phải đi gặp các chuyên gia của chúng, và sau đó trở về nhà.