CODA: Đột phá hiệu năng GPU cho mô hình Transformer nhờ kỹ thuật GEMM-Epilogue
CODA là một trừu tượng hóa kernel GPU mới giúp giải quyết nút thắt hiệu năng trong việc huấn luyện Transformer. Bằng cách viết lại các khối Transformer thành các chương trình GEMM-plus-epilogue, phương pháp này giảm thiểu việc di chuyển dữ liệu giữa bộ nhớ và tăng tốc độ tính toán đáng kể.
Các hệ thống huấn luyện Transformer hiện nay chủ yếu được xây dựng dựa trên nền tảng của đại số tuyến tính mật độ cao (GEMM). Tuy nhiên, một phần không nhỏ thời gian xử lý thực tế lại bị tiêu tốn cho các toán tử phụ thuộc vào bộ nhớ (memory-bound operators) xung quanh.
Các thao tác như chuẩn hóa (normalization), kích hoạt (activations), cập nhật residual và các phép tính giảm dữ liệu (reductions) liên tục di chuyển các tensor trung gian khổng lồ qua bộ nhớ toàn cục (global memory) trong khi thực hiện rất ít phép tính số học. Điều này khiến việc di chuyển dữ liệu trở thành một nút thắt ngày càng quan trọng, ngay cả khi các stack huấn luyện đã được tối ưu hóa rất cao.
Để giải quyết vấn đề này, nhóm nghiên cứu đã giới thiệu CODA — một trừu tượng hóa kernel GPU mới nhằm thể hiện các phép tính này dưới dạng các chương trình GEMM-plus-epilogue.
Nguyên lý hoạt động của CODA
CODA dựa trên quan sát rằng nhiều toán tử trong Transformer, vốn được triển khai dưới dạng các kernel riêng biệt trong các framework, có thể được tham số hóa lại về mặt đại số để thực thi ngay khi ô (tile) đầu ra của GEMM vẫn còn nằm trên chip, trước khi được ghi vào bộ nhớ.
Trừu tượng hóa này giữ nguyên vòng lặp chính của GEMM và mở rộng một tập nhỏ các nguyên thủy epilogue có thể kết hợp cho các tác vụ như:
- Scaling (tỷ lệ)
- Reductions (giảm dữ liệu)
- Pairwise transformations (biến đổi cặp)
- Accumulation (tích lũy)
Giao diện được giới hạn này giúp bảo toàn cấu trúc hiệu năng của các GEMM được viết bởi các chuyên gia, đồng thời vẫn đủ linh hoạt để bao phủ gần như mọi phép tính không phải attention trong quá trình truyền xuôi (forward pass) và truyền ngược (backward pass) của một khối Transformer tiêu chuẩn.
Hiệu suất và Tương lai
Trên các khối lượng công việc đại diện cho Transformer, cả các kernel CODA do con người viết và do LLM (Mô hình ngôn ngữ lớn) tạo ra đều đạt được hiệu suất cao.
Kết quả này gợi ý rằng lập trình theo mô hình GEMM-plus-epilogue là một con đường thực tế để kết hợp năng suất phát triển ở cấp độ framework với hiệu quả tối đa ở cấp độ phần cứng. Điều này đặc biệt có ý nghĩa trong bối cảnh nhu cầu tính toán cho AI ngày càng tăng, đòi hỏi việc tận dụng tối đa phần cứng GPU là bắt buộc.
