Kỹ thuật Gradient Checkpointing đánh đổi tính toán lấy bộ nhớ (Activation Checkpointing)

Mô tả bài toán

Khi huấn luyện mô hình Transformer lớn, phần lớn bộ nhớ VRAM không phải bị chiếm bởi trọng số mô hình mà bởi Activation Cache (các kết quả trung gian cần lưu lại cho pha backward). Kỹ thuật Gradient Checkpointing (Chen et al., 2016, torch.utils.checkpoint):

• Trong Forward pass: chỉ lưu trữ activation tại một số nút đánh dấu ("checkpoints"). Toàn bộ activation trung gian giữa các checkpoint bị hủy bỏ ngay lập tức để giải phóng RAM.
• Trong Backward pass: khi cần gradient của các tầng trung gian, hệ thống chạy lại forward pass cục bộ từ checkpoint gần nhất.

Đánh đổi: Tăng thêm ≈ 25% - 30% thời gian tính toán nhưng giảm tới 70% - 80% dung lượng bộ nhớ đỉnh (Peak Memory), cho phép huấn luyện mô hình gấp 4 lần kích thước thông thường!

Hãy mô phỏng cơ chế này qua hàm simulate_checkpointed_memory(num_layers: int, activation_size_mb: float, checkpoint_every: int) -> dict:

• num_layers: tổng số tầng mạng N.
• activation_size_mb: dung lượng activation của mỗi tầng (MB).
• checkpoint_every: tần suất lưu checkpoint K (cứ mỗi K tầng thì lưu 1 checkpoint).
• Tính:
• memory_standard_mb: dung lượng activation đỉnh trong chế độ chuẩn (lưu tất cả N tầng: N × size).
• memory_checkpointed_mb: dung lượng activation đỉnh khi dùng checkpoint: số checkpoint lưu trữ lceil N / K rceil × size cộng với activation tối đa của một đoạn cục bộ K × size.
• memory_savings_percent: tỷ lệ phần trăm bộ nhớ tiết kiệm được standard - checkpointedstandard × 100%.
• Trả về dictionary chứa 3 giá trị trên.

Ràng buộc

• Thời gian chạy tối đa: 6000ms.
• Giới hạn bộ nhớ: 512MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.