Meta tối ưu hóa GEM: Bí quyết tăng gấp đôi hiệu suất huấn luyện mô hình quảng cáo quy mô LLM

Kiến trúc của GEM và những thách thức huấn luyện đặc thù

GEM là mô hình nền tảng gợi ý trung tâm đứng sau hệ thống quảng cáo của Meta. Mô hình này sở hữu kiến trúc lai với hàng nghìn tỷ tham số embedding thưa (sparse) và hàng tỷ tham số dày (dense). GEM được huấn luyện dựa trên dữ liệu nội dung quảng cáo và tương tác người dùng với hai loại đặc trưng: đặc trưng chuỗi (ví dụ: lịch sử hoạt động của người dùng) và đặc trưng không phải chuỗi (ví dụ: vị trí người dùng, biểu diễn nội dung quảng cáo). Các cơ chế chú ý (attention) tùy chỉnh được áp dụng độc lập cho từng nhóm, đồng thời cho phép học hỏi chéo giữa các đặc trưng.

Sự tương tác giữa kiến trúc lai này và các đặc tính dữ liệu trong lĩnh vực gợi ý (rec-domain) chính là yếu tố khiến việc huấn luyện GEM trở nên khó khăn một cách đặc thù.

Thách thức 1: Đạt hiệu suất sử dụng GPU cao

Các GPU trung tâm dữ liệu và các ngăn xếp phần mềm (software stacks) hiện nay chủ yếu được tối ưu hóa cho các khối lượng công việc LLM, trong khi khối lượng công việc gợi ý lại có cấu hình khác biệt căn bản do các đặc tính dữ liệu độc đáo và các mô hình tương tác phong phú giữa người dùng và quảng cáo. Điều này khiến việc đạt được hiệu suất sử dụng tính toán GPU cao để huấn luyện một mô hình gợi ý nền tảng có quy mô như GEM trở nên cực kỳ khó khăn.

Thách thức 2: Mở rộng quy mô hiệu quả trên hàng nghìn GPU

Việc huấn luyện GEM trên hàng nghìn GPU với hàng nghìn tỷ tham số embedding thưa và hàng tỷ tham số dày đòi hỏi khả năng mở rộng hiệu quả, chứ không chỉ là tăng quy mô đơn thuần. Việc chỉ thêm nhiều GPU không đồng nghĩa với việc tăng tốc tương ứng. Trong huấn luyện phân tán, độ trễ E2E (end-to-end) cho mỗi bước huấn luyện được xác định bởi:

Độ trễ E2E = Giá trị lớn nhất trên các GPU Rank (Max(Thời gian tính toán cục bộ, Thời gian giao tiếp))

Để đạt được khả năng mở rộng gần như tuyến tính, cần có bốn điều kiện:

Khối lượng công việc của GEM đe dọa tất cả các điều kiện này:

Cách tiếp cận và Khung hiệu suất của chúng tôi

Với những thách thức nêu trên, chúng tôi cần một khung làm việc giúp chuyển đổi nỗ lực đồng thiết kế (co-design) dàn trải thành một số ít các đòn bẩy kỹ thuật. Chúng tôi đo lường hiệu suất huấn luyện thông qua E2E MFU, được phân tách thành hai yếu tố:

E2E MFU = Local MFU (hiệu suất tính toán) × Scaling Ratio (hiệu suất mở rộng)

Các yếu tố này mô tả hai vấn đề tối ưu hóa có liên quan nhưng riêng biệt.

Local MFU (hiệu suất tính toán) đo lường mức độ sử dụng các đơn vị tính toán của một GPU đơn lẻ — tức là khối lượng công việc chạy gần với giới hạn phần cứng (hardware roofline) đến mức nào. Chỉ số này được quyết định bởi thiết kế kernel, độ chính xác số học và mức độ phù hợp của các mô hình tính toán (kích thước dữ liệu, độ dài chuỗi) với kiến trúc GPU (Tensor cores, phân cấp bộ nhớ, lập lịch đa xử lý luồng).

Scaling Ratio (hiệu suất mở rộng) đo lường mức độ duy trì hiệu suất của một GPU đơn lẻ khi phân phối trên hàng nghìn GPU. Tỷ lệ mở rộng bằng 1.0 nghĩa là mở rộng tuyến tính hoàn hảo; trong thực tế, chi phí giao tiếp, mất cân bằng tải, hiệu ứng straggler (nút chạy chậm) và việc tính toán lại kích hoạt (activation recomputation) do áp lực bộ nhớ đều làm giảm chỉ số này.

Để tách biệt Local MFU, chúng tôi chạy các lớp mô hình riêng lẻ trên một GPU duy nhất và tính toán MFU trung bình có trọng số mà không cần tính toán lại kích hoạt hoặc lộ ra chi phí giao tiếp. Scaling Ratio được suy ra từ tỷ lệ giữa Local MFU và E2E MFU.

Việc phân tách này rất quan trọng vì nó cho phép chúng tôi xử lý hiệu suất tính toán và hiệu suất mở rộng như những vấn đề tối ưu hóa có liên quan nhưng riêng biệt, mỗi vấn đề đều có bộ kỹ thuật chuyên biệt:

Cả hai đều phải được giải quyết để tối đa hóa E2E MFU.

Tối ưu hóa hiệu suất tính toán với các Recommendation Kernels và huấn luyện độ chính xác cực thấp

Để giải quyết các thách thức đặc thù của hệ thống gợi ý nêu trên và đẩy mạnh hiệu suất sử dụng GPU FLOPS, chúng tôi đã xây dựng một thư viện kernel tùy chỉnh và công thức huấn luyện độ chính xác cực thấp (ultra-low-precision), được thiết kế và tối ưu hóa riêng cho khối lượng công việc gợi ý trên phần cứng GPU mới nhất.

Bên trong Thư viện Kernel tùy chỉnh cho hệ thống gợi ý

Jagged Sequence Flash Attention

FlashAttention được thiết kế cho các chuỗi dày, có độ dài cố định phổ biến trong các LLM. Trong các mô hình gợi ý, chuỗi người dùng vốn dĩ không đồng nhất (jagged) — thay đổi từ hàng trăm đến hàng chục nghìn token trên mỗi mẫu — và việc đệm (padding) đến độ dài tối đa có thể gây lãng phí tới 50% tài nguyên tính toán.

Các triển khai FlashAttention tiêu chuẩn giả định độ dài chuỗi đồng nhất để xếp gạch (tiling) và song song hóa hiệu quả; với đầu vào không đồng nhất, các cách tiếp cận thông thường sẽ phải đệm (gây lãng phí tính toán) hoặc để các SM (Streaming Multiprocessors) nhàn rỗi khi các chuỗi ngắn kết thúc sớm. Chúng tôi đã phát triển JFA, một triển khai FlashAttention tùy chỉnh hoạt động trực tiếp trên các tensor không đồng nhất có độ dài thay đổi, loại bỏ chi phí đệm trong khi vẫn hỗ trợ các tính năng đặc thù của hệ thống gợi ý như bias chú ý tùy chỉnh, độ dài query/key-value bất đối xứng và các bước lan truyền ngược (backward passes) hiệu quả.

Chúng tôi đã phát triển JFA qua bốn thế hệ, dần dần thu hẹp khoảng cách từ việc chậm hơn so với SDPA (scaled dot-product attention) có đệm đến việc đạt hiệu suất ngang bằng với SOTA CUDA/Cutlass trên các GPU thế hệ mới nhất:

JFA v4 (TLX) đạt mức cải thiện 40-140% TFLOPS so với JFA v2, mang lại mức tăng ổn định trong các phân phối không đồng nhất thực tế (độ thưa 0.5), đóng góp vào mức tăng 18.5% Local MFU tương đối và 12% QPS.

Generalized Dot-Product Attention (GDPA)

GEM sử dụng các mô hình tương tác dạng chú ý đa dạng — self-attention, PMA và cross-attention — chia sẻ một cấu trúc chung: hai phép nhân ma trận với một phép kích hoạt theo phần tử ở giữa, nhưng thay thế softmax bằng các hàm kích hoạt như GELU hoặc SiLU. Chúng tôi hợp nhất các mô-đun này dưới một kernel GDPA duy nhất được tối ưu hóa cho khối lượng công việc huấn luyện RecSys thực tế trên các GPU thế hệ mới nhất.

Các kernel FlashAttention hiện có được thiết kế cho đầu vào chuỗi dài, dày kiểu LLM và hoạt động kém hiệu quả dưới lưu lượng thực tế. Chúng tôi quan sát thấy khoảng cách hiệu suất 2.6x trong quá trình lan truyền xuôi (forward) và lên tới 4x trong trường hợp xấu nhất giữa khối lượng công việc thực tế và các benchmark tổng hợp do các chuỗi K/V ngắn/bất đối xứng, đầu vào không đồng nhất và kích thước batch lớn phá vỡ các giả định về chiếm dụng pipeline.

Chúng tôi đã thiết kế lại pipeline kernel, lập lịch và toán học để thu hẹp khoảng cách hiệu suất giữa lưu lượng thực tế và giới hạn phần cứng.

Với những tối ưu hóa này, kernel GDPA đạt tốc độ lan truyền xuôi nhanh gấp 2 lần (1,145 BF16 TFLOPs, ~97% hiệu suất Tensor Core) và nhanh gấp 1.6 lần trong quá trình lan truyền ngược so với baseline. Trong các thiết lập sản xuất với K/V ngắn, nó đạt tốc độ lan truyền xuôi nhanh gấp 3.5 lần so với Flash Attention 4 (FA4). Khi áp dụng trên toàn bộ mô hình, các kernel này mang lại mức cải thiện thông lượng huấn luyện E2E hơn 30%.

BlockAttention

Đối với self-attention của GEM, thách thức hiệu suất cốt lõi là mở rộng các chuỗi người dùng dài mà không phải trả chi phí bậc hai của full attention. Trước tiên, chúng tôi chuyển lớp này từ full self-attention sang sliding-window attention, giới hạn mỗi token với các sự kiện lân cận và giảm độ phức tạp từ O(L2) xuống O(L * window). Điều này làm cho các chuỗi dài trở nên khả thi. Kernel Sliding Window Attention (SWA) đã bỏ qua các ô ngoài cửa sổ trong JFA và giảm độ trễ self-attention cho chuỗi dài tới 68% với NE (normalized entropy, một chỉ số chất lượng mô hình) không đổi.