Khi huấn luyện một mô hình học sâu với 100 lớp, GPU phải lưu toàn bộ giá trị kích hoạt trung gian của từng lớp ở lượt lan truyền thuận, vì lượt lan truyền ngược cần chúng để tính gradient. Với mô hình lớn, bộ nhớ GPU có thể cạn trước khi kịp tăng kích thước batch hoặc thêm lớp. Gradient checkpointing giải quyết bài toán này bằng cách chỉ lưu một số giá trị kích hoạt đã chọn, phần còn lại sẽ được tính lại khi cần.

Bài toán bộ nhớ cho activation

Trong lượt lan truyền thuận thông thường, mỗi lớp tạo ra một tensor kích hoạt. Giả sử mô hình có n lớp, hệ thống phải giữ n tensor trong bộ nhớ cho đến khi lượt lan truyền ngược hoàn tất. Với mô hình ngôn ngữ lớn hiện đại có hàng trăm lớp, mỗi tensor có thể chiếm hàng GB. Tổng bộ nhớ cho activation tăng theo O(n), và đây thường là phần chiếm nhiều VRAM nhất, vượt cả trọng số mô hình.

Chiến lược chia checkpoint theo đoạn

Thay vì lưu toàn bộ, mô hình được chia thành nhiều đoạn và chỉ lưu activation ở ranh giới giữa các đoạn, gọi là checkpoint. Ví dụ, mô hình 16 lớp được chia thành 4 đoạn, mỗi đoạn gồm 4 lớp. Hệ thống chỉ lưu activation tại các lớp 4, 8, 12 và 16. Khi lượt lan truyền ngược cần activation của lớp 6, hệ thống sẽ chạy lại lượt lan truyền thuận từ checkpoint gần nhất, tức lớp 4, qua lớp 5 và 6 để khôi phục giá trị cần thiết.

Tính lại trong lượt lan truyền ngược

Cơ chế hoạt động như sau: lượt lan truyền ngược đi từ lớp cuối trở lên. Khi đến một đoạn, hệ thống nạp activation tại checkpoint ở đầu đoạn, chạy lại lượt lan truyền thuận qua toàn bộ đoạn đó để khôi phục các activation trung gian, rồi mới tính gradient. Sau khi xử lý xong đoạn đó, các activation vừa được tính lại sẽ được giải phóng ngay. Nhờ vậy, tại mỗi thời điểm hệ thống chỉ cần giữ activation của một đoạn trong bộ nhớ.

Cách ước tính mức tiết kiệm bộ nhớ

Nếu mô hình có n lớp và đặt k checkpoint, bộ nhớ cho activation sẽ giảm từ O(n) xuống O(k + n/k). Khi tối thiểu hóa biểu thức này bằng đạo hàm, ta thu được giá trị k tối ưu là sqrt(n). Với mô hình 100 lớp, nếu đặt khoảng 10 checkpoint, bộ nhớ sẽ giảm từ O(100) xuống O(10 + 10) = O(20), tương đương tiết kiệm khoảng 80%. Đổi lại, hệ thống phải chạy thêm một lượt lan truyền thuận cho mỗi đoạn, khiến tổng chi phí tính toán tăng khoảng 33%.

Vì sao sqrt(n) là mức checkpoint tối ưu

Con số sqrt(n) không phải lựa chọn tùy ý. Đây là điểm cân bằng về mặt toán học: nếu đặt ít checkpoint hơn sqrt(n), mỗi đoạn sẽ quá dài, nên bộ nhớ tạm cho đoạn đó vẫn lớn. Nếu đặt nhiều hơn, số checkpoint phải lưu tăng lên nhưng phần bộ nhớ tiết kiệm thêm không đáng kể. Trong thực tế, PyTorch cung cấp torch.utils.checkpoint cho phép bạn chỉ định chính xác layer nào làm checkpoint. Với các framework hiện đại, việc bật gradient checkpointing thường chỉ cần một dòng code, nhưng hiểu nguyên lý giúp bạn chọn vị trí checkpoint tối ưu cho kiến trúc cụ thể.

Một số lưu ý khi áp dụng

Gradient checkpointing không hoàn toàn tương thích với dropout hoặc batch normalization, vì khi chạy lại lượt lan truyền thuận, các lớp này có thể cho kết quả khác lần đầu do mặt nạ ngẫu nhiên hoặc thống kê đang chạy thay đổi. Giải pháp phổ biến là lưu seed ngẫu nhiên hoặc cố định trạng thái khi tính lại. Kỹ thuật này đã trở thành lựa chọn phổ biến khi fine-tune LLM trên một GPU, thường được kết hợp với LoRA và mixed precision để tối đa hóa kích thước mô hình có thể huấn luyện.