Mô hình AI mạnh nhất nào có thể được huấn luyện trên laptop chỉ trong 5 phút?

• Huấn luyện một mô hình transformer kiểu GPT khoảng 1.8M tham số trên MacBook Pro trong 5 phút với khoảng 20M token TinyStories, đạt khoảng 9.6 perplexity
• Ràng buộc chính của việc huấn luyện trong 5 phút là kích thước mô hình và số token có thể xử lý; mô hình càng lớn thì hội tụ càng chậm và hiệu quả giảm do dữ liệu ít
• Về tối ưu hiệu năng, dùng MPS, còn thay vì compile/lượng tử hóa/tích lũy gradient hay thay PyTorch thì chọn mô hình nhỏ hơn hiệu quả hơn
• Các bộ dữ liệu đơn giản và nhất quán như TinyStories có tác động tích cực hơn đến hiệu năng của mô hình nhỏ so với dữ liệu kiểu bách khoa toàn thư
• Kiến trúc transformer cho kết quả tốt hơn LSTM hay diffusion trong điều kiện kích thước nhỏ và thời gian huấn luyện ngắn

---

Tổng quan

Bài viết này trình bày kết quả thử nghiệm về mô hình ngôn ngữ AI có hiệu năng tối đa có thể được huấn luyện trên laptop (MacBook Pro) trong 5 phút, cùng với các insight về chiến lược huấn luyện tối ưu, lựa chọn bộ dữ liệu và kiến trúc mô hình.

Tóm tắt kết quả thử nghiệm

• Huấn luyện một mô hình transformer kiểu GPT khoảng 1.8M tham số với khoảng 20M dữ liệu TinyStories, ghi nhận 9.6 perplexity
• Ví dụ sinh văn bản tuy ngắn nhưng có dạng câu chuyện nhất quán, với ngữ pháp tiếng Anh nhìn chung được giữ đúng
• Nhấn mạnh rằng kết quả mô hình ở mức thực dụng trong vòng 5 phút là tốt hơn kỳ vọng

Bối cảnh và giới hạn của thử nghiệm

• Đây là một thử nghiệm xuất phát từ sự tò mò mang tính không thực tế: huấn luyện nhanh một mô hình mạnh trong môi trường laptop
• Trên thực tế có thể huấn luyện các mô hình mạnh hơn bằng GPU hiệu năng cao trên cloud (như H100), nhưng ràng buộc của thử nghiệm này là thời gian (5 phút)
• Kích thước mô hình càng lớn thì tốc độ xử lý token càng chậm, khiến việc đạt kết quả tốt trong 5 phút trở nên khó khăn
  • Các mô hình quá nhỏ (ví dụ: 10K tham số) không thể học đủ độ phức tạp
  • Phạm vi thực dụng là các mô hình khoảng 1M~2M tham số

Tối ưu thông lượng xử lý

• Dùng MPS (Metal Performance Shaders của Apple) là hiệu quả nhất
• Các tối ưu toán học như torch.compile, float16, MLX... cho hiệu quả cải thiện thấp hơn kỳ vọng, thậm chí còn làm giảm hiệu năng
• Tích lũy gradient có mục đích quản lý bộ nhớ, nhưng trên thực tế làm giảm tốc độ nghiêm trọng
• Mô hình cần có khả năng cập nhật weight nhanh trong bộ nhớ nội bộ để đạt hiệu quả

Lựa chọn bộ dữ liệu

• Khi dùng số token hạn chế (khoảng 10~20M), ban đầu tác giả thử dữ liệu wiki tiếng Anh đơn giản như Simple English Wikipedia; kết quả là giữ được tính nhất quán ngữ pháp nhưng thiếu nhất quán về ngữ nghĩa
  • Nội dung thiên về danh từ riêng, liệt kê các sự kiện có cảm giác gượng ép nên bị hạn chế trong việc tạo ra nội dung có ý nghĩa
• Khi dùng bộ dữ liệu TinyStories, do cấu trúc câu chuyện rõ ràng và ngôn ngữ đơn giản nên kết quả nhất quán và có ý nghĩa hơn nhiều
  • Đây là các câu chuyện ở mức trẻ 4 tuổi, nên mô hình nhỏ cũng học tốt

Tokenizer và token hóa

• Việc huấn luyện tokenizer không được tính trong 5 phút, và do quy mô dữ liệu nhỏ nên nhu cầu tối ưu hóa cũng thấp
• Học các token đa byte dễ hơn cho quá trình huấn luyện mô hình

Thử nghiệm kiến trúc mô hình

• Sử dụng kiến trúc transformer (kiểu GPT-2)

  • Điều chỉnh các siêu tham số như 2~3 layer, hàm kích hoạt SwiGLU, positional embedding...
  • LSTM có hiệu năng gần sát nhưng transformer vẫn tốt hơn về perplexity
  • Dropout, mixture-of-experts... không hiệu quả vì quy mô quá nhỏ
  • Curriculum learning gần như không có tác dụng vì thời gian huấn luyện quá ngắn

• Thử mô hình diffusion (D3PM)

  • Vì ngôn ngữ tự nhiên là các token rời rạc, quá trình khuếch tán chỉ tạo ra các token ngẫu nhiên vô nghĩa nên thất bại
  • So với transformer hay LSTM, rất khó hình thành cấu trúc câu nhanh

Mối quan hệ giữa kích thước mô hình và thông lượng token/giây

• Các mô hình 1M~2M tham số là điểm ngọt lý tưởng nhất
  • Quá lớn thì không thể hội tụ trong 5 phút, quá nhỏ thì chạm trần hiệu năng gần như ngay lập tức
• Kết quả thử nghiệm nhìn chung phù hợp với Chinchilla scaling law
  • Kích thước mô hình lý tưởng là tổng số token huấn luyện/20, và điều này cũng được xác nhận trong thử nghiệm này

Kết luận và hàm ý

• Ngay cả với thời gian rất ngắn và phần cứng nhỏ, việc huấn luyện một mô hình kể chuyện nhất quán vẫn khả thi
• Việc huấn luyện trong 5 phút không phù hợp để phát triển mô hình mạnh, nhưng vẫn có ý nghĩa với thiết kế mô hình nhỏ, siêu nhẹ và các thử nghiệm tối ưu phần cứng/kiến trúc
• Trong tương lai, khi GPU trên laptop và cấu trúc mô hình tiếp tục phát triển, hiệu năng của các mô hình có thể được huấn luyện chỉ trong vài phút vẫn còn nhiều tiềm năng cải thiện