MegaTrain: Huấn luyện mô hình ngôn ngữ lớn 100 tỷ tham số trên một GPU duy nhất

MegaTrain: Huấn luyện mô hình ngôn ngữ lớn 100 tỷ tham số trên một GPU duy nhất

Việc huấn luyện các mô hình ngôn ngữ lớn (LLM) thường đòi hỏi hàng nghìn GPU và chi phí tính toán cực kỳ lớn, tạo ra rào cản lớn đối với nhiều tổ chức và nhà nghiên cứu. Tuy nhiên, một nghiên cứu mới mang tên MegaTrain đã mở ra hướng đi hoàn toàn khác, cho phép chạy các mô hình hơn 100 tỷ tham số trên một GPU duy nhất mà không cần giảm độ chính xác.

Cách tiếp cận lấy bộ nhớ làm trung tâm

Khác với các hệ thống truyền thống tập trung vào GPU, MegaTrain hoạt động như một hệ thống lấy bộ nhớ làm trung tâm (memory-centric). Nó coi GPU chỉ là một động cơ tính toán tạm thời, trong khi toàn bộ tham số của mô hình và trạng thái của bộ tối ưu hóa (optimizer states) được lưu trữ trong bộ nhớ máy chủ (CPU RAM).

Quá trình huấn luyện diễn ra theo cơ chế luân chuyển: đối với mỗi lớp, hệ thống sẽ truyền tham số vào GPU để tính toán và đẩy độ dốc (gradients) ra ngoài ngay lập tức. Cách tiếp cận này giúp giảm thiểu tối đa lượng trạng thái cần lưu trữ cố định trên thiết bị, cho phép xử lý các mô hình có kích thước vượt xa khả năng lưu trữ của GPU.

Các tối ưu hóa kỹ thuật then chốt

Để giải quyết vấn đề nút thắt băng thông giữa CPU và GPU, nhóm nghiên cứu đã áp dụng hai giải pháp tối ưu hóa quan trọng:

• Động cơ thực hiện có đường ống và bộ đệm kép: Kỹ thuật này chồng chéo các tác vụ lấy trước tham số (prefetching), tính toán và đẩy dữ liệu (offloading) trên nhiều luồng CUDA khác nhau. Điều này cho phép GPU hoạt động liên tục không ngừng nghỉ, tận dụng tối đa khả năng tính toán thay vì phải chờ đợi dữ liệu.
• Mẫu lớp không trạng thái: Thay thế các biểu đồ autograd cố định bằng các mẫu lớp linh hoạt. Các trọng số được liên kết động khi dữ liệu được luân chuyển vào, giúp loại bỏ metadata thừa thãi của biểu đồ đồ thị đồng thời cung cấp sự linh hoạt cao hơn trong việc lập lịch tác vụ.

Hiệu năng ấn tượng

Trên cấu hình sử dụng một GPU H200 kết hợp với 1,5TB bộ nhớ RAM, MegaTrain đã chứng minh khả năng huấn luyện đáng tin cậy các mô hình lên tới 120 tỷ tham số.

Hơn nữa, khi huấn luyện mô hình 14 tỷ tham số, hệ thống này đạt tốc độ xử lý (throughput) cao gấp 1,84 lần so với DeepSpeed ZeRO-3 khi sử dụng cơ chế offloading CPU. Đáng chú ý, MegaTrain còn cho phép huấn luyện mô hình 7 tỷ tham số với độ dài ngữ cảnh khổng lồ lên tới 512k token trên một GPU GH200 duy nhất.

Đây là bước tiến quan trọng, giúp dân hóa việc huấn luyện các mô hình AI siêu lớn, cho phép các nhà nghiên cứu tiếp cận công nghệ này mà không cần đầu tư vào hạ tầng dữ liệu quy mô lớn.