Tóm tắt
MegaTrain là một hệ thống tập trung vào bộ nhớ, cho phép huấn luyện mô hình ngôn ngữ lớn với hơn 100B tham số ở độ chính xác đầy đủ chỉ trên một GPU duy nhất. Không giống như hệ thống tập trung vào GPU truyền thống, MegaTrain lưu trữ các tham số và trạng thái tối ưu hóa trong bộ nhớ CPU và coi GPU như động cơ tính toán tạm thời.
Phát hiện chính
- MegaTrain truyền tham số cho mỗi lớp và xuất gradient để giảm thiểu trạng thái thiết bị kéo dài.
- Hai tối ưu hóa chính được áp dụng: bội gối nhiệm vụ, thực hiện truy cập kép lồng ghép, và thay thế đồ thị autograd bền bỉ bằng mẫu lớp vô hình.
- Hệ thống có thể huấn luyện mô hình lên đến 120B tham số trên một GPU H200 với 1.5TB bộ nhớ máy chủ.
- Nó tiết kiệm 1.84 lần thông lượng so với DeepSpeed ZeRO-3 khi huấn luyện 14B mô hình.
Ý nghĩa với Dev
- MegaTrain mang lại cơ hội xử lý mô hình ngôn ngữ lớn hơn mà không cần đầu tư vào nhiều phần cứng GPU đắt tiền.
- Dev có thể tận dụng hệ thống này để huấn luyện mô hình với bối cảnh 512k token trên một GH200 duy nhất, tối ưu hóa tài nguyên mà vẫn đạt kết quả tốt.
- Điều này rất hữu ích trong việc phát triển sản phẩm AI hoặc mở rộng quy mô nghiên cứu với chi phí thấp hơn.