Tối ưu hóa hiệu suất Deep Learning từ những nguyên lý cơ bản

Bạn muốn cải thiện hiệu suất của mô hình Deep Learning, nhưng không biết nên bắt đầu từ đâu? Thay vì áp dụng một loạt các mẹo vặt vô thưởng vô phạt được nghe thấy trên mạng, cách tiếp cận khoa học hơn là quay về với những nguyên lý cơ bản (first principles).

GPU hoạt động hết công suất

Hiệu suất của hệ thống Deep Learning thực chất phụ thuộc vào ba thành phần chính: tính toán (compute), băng thông bộ nhớ (memory bandwidth) và chi phí quản lý (overhead). Việc xác định xem hệ thống của bạn đang bị kẹt ở đâu sẽ giúp bạn tìm ra giải pháp tối ưu đúng đắn nhất.

Tối đa hóa thời gian tính toán (Compute-bound)

Mục tiêu tối thượng của chúng ta là giữ cho GPU luôn ở trạng thái "compute-bound" (bị giới hạn bởi khả năng tính toán). Bạn đã trả tiền cho sức mạnh khủng khiếp của GPU (ví dụ: 312 teraflops trên Nvidia A100), vì vậy bạn muốn tận dụng tối đa con số đó.

Thông số kỹ thuật GPU

Tuy nhiên, các bộ tăng tốc Machine Learning hiện đại đều có phần cứng chuyên dụng cho phép nhân ma trận (Tensor Cores). Nếu bạn không thực hiện nhân ma trận, hiệu suất thực tế sẽ giảm đi rất nhiều. May mắn là trong các mô hình như BERT, các phép toán không phải nhân ma trận (như chuẩn hóa hoặc hàm kích hoạt) chỉ chiếm một phần rất nhỏ tổng số FLOPS (dưới 0.2%). Do đó, chúng không phải là vấn đề về mặt tính toán.

Vấn đề về băng thông bộ nhớ (Memory-bandwidth bound)

Nếu các phép toán không phải nhân ma trận tốn rất ít tài nguyên tính toán, tại sao chúng đôi khi lại làm chậm hệ thống? Thủ phạm chính là việc di chuyển dữ liệu.

Hãy tưởng tượng GPU như một nhà máy. Nhà máy (các đơn vị tính toán) rất hiệu quả nhưng không có nhiều không gian lưu trữ. Kho hàng (DRAM - bộ nhớ GPU) thì rộng lớn nhưng nằm ở xa. Chi phí vận chuyển nguyên liệu từ kho về nhà máy chính là "memory bandwidth cost".

Khi thực hiện một phép toán đơn giản như torch.cos, chúng ta tốn rất nhiều thời gian để chuyển dữ liệu từ DRAM lên đơn vị tính toán, thực hiện một phép tính nhỏ xíu, rồi chuyển ngược lại. Kết quả là phần lớn thời gian bị lãng phí cho việc vận chuyển chứ không phải sản xuất. Đây được gọi là trạng thái "memory-bound".

Mô hình nhà máy và kho hàng

Giải pháp: Operator Fusion

Để giải quyết vấn đề trên, kỹ thuật quan trọng nhất trong các trình biên dịch Deep Learning hiện nay là Operator Fusion (gộp toán tử).

Thay vì thực hiện từng toán tử riêng lẻ và phải đọc/ghi dữ liệu liên tục ra vào bộ nhớ chính (DRAM), chúng ta "gộp" chúng lại. Giữ dữ liệu lại tại nhà máy (nhớ nhanh/registers), thực hiện chuỗi các phép tính, rồi mới ghi kết quả cuối cùng về DRAM một lần.

Ví dụ, x.cos().cos() nếu không gộp sẽ cần 4 lần truy cập bộ nhớ toàn cục. Nếu gộp lại, chỉ cần 2 lần. Điều này giúp tăng tốc gấp đôi. Các công cụ như NVFuser, XLA hoặc viết kernel tùy chỉnh bằng Triton đều giúp thực hiện việc này.

Chi phí quản lý (Overhead)

Ngoài tính toán và bộ nhớ, còn một kẻ thù nữa: Overhead. Đây là thời gian dành cho các việc không phải tính toán hay chuyển dữ liệu, như chạy trình thông dịch Python, khởi chạy kernel CUDA, hoặc các lớp dispatch của PyTorch.

Python cực kỳ chậm so với GPU. Trong thời gian Python thực hiện một phép cộng, GPU A100 có thể thực hiện gần 10 triệu phép tính. Tuy nhiên, các framework như PyTorch thực thi bất đồng bộ (asynchronously), cho phép CPU xếp hàng các lệnh cho GPU trong khi GPU đang bận. Điều này giúp ẩn đi phần lớn chi phí overhead.

Để biết hệ thống có bị dính lỗi overhead không, hãy thử tăng kích thước dữ liệu (batch size). Nếu thời gian chạy không tăng tương ứng, bạn đang bị giới hạn bởi overhead.

Kết luận

Để tối ưu hóa hiệu suất Deep Learning, bước quan trọng nhất là xác định nút thắt cổ chai (bottleneck) của mô hình:

1. Nếu bạn đạt gần 100% FLOPS tối đa của GPU: Bạn đang ở trạng thái compute-bound. Tốt rồi, hãy giữ nguyên.
2. Nếu FLOPS thấp nhưng băng thông bộ nhớ cao: Bạn đang memory-bound. Hãy cân nhắc operator fusion hoặc giảm thiểu truy cập bộ nhớ.
3. Nếu tăng kích thước batch mà thời gian không đổi: Bạn đang overhead-bound. Hãy giảm thiểu các thao tác của Python hoặc sử dụng tracing/JIT.

Hiểu rõ các nguyên lý này sẽ giúp bạn đưa ra quyết định chính xác hơn là chỉ áp dụng các mẹo vặt một cách mù quáng.