Khắc phục lỗi OOM trong PyTorch: Huấn luyện đồ thị khổng lồ trên GPU thông thường

Bài viết giới thiệu Disk Sparse Adam (DSA), một trình tối ưu hóa ngoài bộ nhớ mới dành cho PyTorch, được thiết kế để huấn luyện các mô hình thưa thớt quy mô lớn, chẳng hạn như các embedding đồ thị tri thức với hàng triệu nút, trên các GPU cấp tiêu dùng. SparseAdam chuẩn của torch.optim lưu trữ trạng thái mô-men bậc nhất và bậc hai trong bộ nhớ, điều này đối với mô hình có 10 triệu thực thể và vector 128 chiều sẽ tiêu tốn khoảng 10,24 GB bộ nhớ bổ sung, thường dẫn đến lỗi hết bộ nhớ CUDA. Thay vào đó, DSA lưu trữ các ma trận mô-men này trên đĩa dưới dạng tệp nhị phân và ánh xạ chúng vào bộ nhớ bằng cơ chế mmap của hệ điều hành, chỉ đọc và ghi trạng thái của batch đang hoạt động trên mỗi bước tối ưu hóa. Trình tối ưu hóa được thiết kế như một giải pháp thay thế tương thích trực tiếp cho các pipeline PyTorch tiêu chuẩn và hỗ trợ cả không gian Euclid và hyperbolic (quả cầu Poincaré). Tác giả báo cáo rằng mức sử dụng bộ nhớ cho trạng thái trình tối ưu hóa giảm từ gigabyte xuống còn megabyte, và một điểm chuẩn trên Kaggle với 1 triệu thực thể cho thấy mức sử dụng VRAM là 0,00 MB và tốc độ xử lý đạt 134.212 mẫu mỗi giây. DSA có sẵn trên GitHub theo giấy phép MIT, với các hạn chế bao gồm yêu cầu ổ SSD NVMe để có hiệu suất tốt nhất và chỉ tương thích với gradient thưa thớt, chẳng hạn như từ torch.nn.Embedding hoặc EmbeddingBag.