MegaTrain: Đào tạo chính xác toàn bộ hơn 100 tỷ LLM thông số trên một GPU đơn

Xem PDF HTML (thử nghiệm)

Tóm tắt:Chúng tôi giới thiệu MegaTrain, một hệ thống lấy bộ nhớ làm trung tâm giúp huấn luyện hiệu quả các mô hình ngôn ngữ lớn với hơn 100 tỷ tham số ở độ chính xác đầy đủ trên một GPU duy nhất. Không giống như các hệ thống tập trung vào GPU truyền thống, MegaTrain lưu trữ tham số và trạng thái bộ tối ưu hóa trong bộ nhớ máy chủ (bộ nhớ CPU) và coi các GPU là các công cụ tính toán tạm thời. Đối với mỗi lớp, chúng tôi truyền phát tham số vào và xuất gradient ra, giúp giảm thiểu trạng thái thiết bị tồn tại lâu dài. Để giải quyết nút thắt băng thông CPU-GPU, chúng tôi áp dụng hai tối ưu hóa chính. 1) Chúng tôi giới thiệu một công cụ thực thi đệm kép theo đường ống (pipelined double-buffered) giúp chồng lấp quá trình tìm nạp trước tham số, tính toán và giảm tải gradient trên nhiều luồng CUDA, cho phép GPU thực thi liên tục. 2) Chúng tôi thay thế các biểu đồ autograd cố định bằng các mẫu lớp không trạng thái (stateless), liên kết trọng số một cách linh hoạt khi chúng được truyền vào, loại bỏ siêu dữ liệu biểu đồ cố định đồng thời mang lại sự linh hoạt trong việc lập lịch. Trên một GPU H200 duy nhất với 1.5TB bộ nhớ máy chủ, MegaTrain huấn luyện ổn định các mô hình lên đến 120 tỷ tham số. Nó cũng đạt được tốc độ huấn luyện gấp 1.84 lần so với DeepSpeed ZeRO-3 với tính năng giảm tải CPU khi huấn luyện các mô hình 14 tỷ tham số. MegaTrain cũng cho phép huấn luyện mô hình 7 tỷ tham số với bối cảnh 512k token trên một thiết bị GH200 duy nhất.

Lịch sử gửi bài

Từ: Zhengqing Yuan [xem email]
[v1] Thứ Hai, 6 tháng 4 năm 2026 18:43:56 UTC (787 KB)