Hugo Vergnes huấn luyện mô hình 3.8B với chi phí dưới 1.000 USD

Nằm ở khoảng giữa "đồ chơi nanoGPT" và "cần một phòng thí nghiệm nghiên cứu", có một vùng rộng lớn, chưa được mô tả kỹ lưỡng, nơi một cá nhân với vài nghìn đô la có thể huấn luyện một mô hình có ý nghĩa.

Tôi muốn tự mình chứng kiến ngôn ngữ và sự hiểu biết nảy sinh từ các trọng số ngẫu nhiên, đồng thời học những phần mà bạn chỉ có thể học được khi bắt đầu từ con số không. Dự án này được viết vào các buổi tối, gỡ lỗi trên một chiếc 5090 và hoàn thiện trên các máy B200 thuê ngoài. Nó được lấy cảm hứng mạnh mẽ từ nanochat của Andrej Karpathy.

Kết quả là một mô hình 3,8 tỷ tham số đạt 0,384 điểm trên CORE, được huấn luyện trên 65 tỷ token trong 43 giờ với chi phí 998 đô la.

Những gì theo sau đây là những gì đã hiệu quả, những gì không, và những gì tôi vẫn chưa biết.

Mô hình của tôi lớn hơn nanochat d32 và mất thời gian thực tương đương. B200 mang lại giá trị tốt hơn trên mỗi đơn vị công việc so với H100. Nhưng với số tiền gần như tương đương với cấu hình 1.000 đô la của nanochat, mô hình này vượt trội hơn hẳn. Đây là một dữ liệu đáng khích lệ về những gì có thể đạt được bên ngoài một phòng thí nghiệm hoặc một công ty lớn với ngân sách tính toán hàng triệu đô la. Khi ranh giới công nghệ dịch chuyển, 1.000 đô la sẽ đưa bạn đi xa hơn nữa.

Thiết lập

Tôi đã xây dựng little-lm như một framework dựa trên cấu hình để huấn luyện các LLM nhỏ chỉ dùng decoder. Mỗi lần chạy được chỉ định đầy đủ bởi một tệp YAML: mô hình, tập dữ liệu, bộ tối ưu hóa (optimizer), lịch trình, các callback. Các thành phần tự đăng ký vào một registry toàn cục và được phân giải theo tên, vì vậy việc thay đổi bộ tối ưu hóa hoặc tập dữ liệu chỉ là thay đổi một dòng trong cấu hình.

Cơ sở hạ tầng tốt sẽ tự hoàn vốn gần như ngay lập tức. Kỷ luật kỹ thuật phần mềm thông thường (như phân tách các mối quan tâm, giao diện sạch, các thành phần có thể thay thế) rất quan trọng trong công việc AI. Nó tiêu tốn của tôi một chút thời gian lúc đầu, và thêm vài lần nữa sau đó để sửa các hợp đồng kém hoặc các điểm chưa tối ưu. Nhưng khoản đầu tư thời gian này sẽ tự hoàn vốn ngay khi bạn gặp vấn đề hội tụ đầu tiên. Tôi nhận thấy rằng cơ sở hạ tầng tuyệt vời là cơ sở hạ tầng gần như không bao giờ yêu cầu bạn phải chỉnh sửa mã nguồn thủ công. Nếu bạn có thể đọc cấu hình và hiểu chính xác những gì xảy ra, và không có cơ chế ẩn nào, điều đó có nghĩa là bạn đã làm tốt. Báo cáo sau đây là kết quả của việc có thể biểu diễn các thử nghiệm dưới dạng một file YAML diff ba dòng thay vì một nhánh (branch).

Mô hình cuối cùng mang phong cách Llama: RMSNorm, RoPE, GQA (24 query head, 8 KV head), relu² MLP, QK-norm, logit softcap, các scalar residual có thể học được theo từng lớp, và các value embedding kiểu ResFormer.

Cần lưu ý rằng các value embedding chiếm 19% số lượng tham số. 14 bảng từ vựng × kv_dim, xuất hiện ở mỗi lớp xen kẽ.

Kết quả

Các thử nghiệm ban đầu

Trước khi có những lần chạy tốt, đã có rất nhiều lần chạy thất bại.

Tôi đã huấn luyện một mô hình Llama 858M trên FineWeb-Edu với 16,4 tỷ token, mất 5,8 ngày trên một chiếc A100 duy nhất. Sử dụng AdamW ở mức 2,5e-4, cosine decay về 0, 5% warmup, batch 256 thông qua gradient accumulation, context 2048.

Kết quả: PIQA đạt 60,45%. GPT-2 124M đạt khoảng 63%. Tôi đã dành sáu ngày tính toán để xây dựng một thứ tồi tệ hơn cả một mô hình nhỏ hơn gấp bảy lần từ năm 2019. Các văn bản tạo ra bị lặp lại và gần như vô nghĩa.

Đường cong mất mát (loss curve) đã nói lên tất cả.

Năm thay đổi đã được rút ra sau quá trình phân tích đó. Tổng hợp lại, chúng tạo ra sự khác biệt giữa lần chạy trên và một mô hình đánh bại GPT-2 với khoảng cách lớn.

Lịch trình LR hình thang (Trapezoidal LR schedule). Warmup 5%. Giữ phẳng và kết thúc với việc giảm dần tuyến tính trong 50% đến 5% cuối cùng của đỉnh. Điểm mấu chốt là mô hình tiếp tục học cho đến cuối thay vì "trôi" qua giai đoạn cuối. Trong lần chạy 3,8B, eval loss vẫn đang giảm ở bước cuối cùng, đó chính xác là hành vi mà lần chạy 858M đã không làm được.

Muon cho các tham số ma trận, AdamW cho mọi thứ khác. Muon chậm hơn trên mỗi bước (phép trực giao hóa Newton-Schulz không miễn phí, mất khoảng 25% trong benchmark tích lũy nông) nhưng chi phí đó chỉ trả một lần cho mỗi bước tối ưu hóa: với 7 bước gradient-accumulation, nó giảm xuống còn ~4%. Nếu đo theo tổng thời gian chạy, sự hội tụ nhanh hơn nhiều.

ClimbMix thay vì FineWeb-Edu. Đây là một bước nhảy vọt về tốc độ hội tụ. Chính xác như những gì Karpathy đã phát hiện.

FP8 + đệm từ vựng (vocab padding). Huấn luyện FP8 thông qua torch._scaled_mm với tính năng dynamic tensorwise scaling trên cả ba GEMM, và đệm từ vựng từ 50.257 lên 50.304 (bội số của 64) để các tensor core hoạt động hiệu quả. Tổng cộng, tăng 33% thông lượng, chủ yếu từ fp8.

Context 1024 thay vì 2048. Giảm một nửa context giúp tăng gấp đôi batch size ở mức bộ nhớ cố định. Thông lượng hầu như không thay đổi trên mỗi token. Chúng ta vẫn bị chi phối bởi các MLP, đó là dấu hiệu tốt cho thấy chúng ta đang sử dụng phần cứng hiệu quả. Dưới đây chúng ta sẽ thảo luận về tác động của độ dài context đối với mô hình.

Đây là toàn bộ quá trình chạy:

~480.000 token/giây ở trạng thái ổn định, đạt 57,3 tỷ token trong 33 giờ. Thời gian thực là 35,9 giờ. Sự khác biệt nằm ở các đánh giá CORE, mỗi lần mất khoảng 15 phút (mười lần trong suốt quá trình chạy) và tiêu tốn 7% tổng thời gian.

Chạy lại công thức tương tự này ở context 2048 token đạt 0,3840. Hầu như toàn bộ khoảng cách đó hóa ra là do một số tác vụ phụ thuộc rất nhiều vào context.

Về bản thân các GPU: 92% hoạt động SM, 40% chiếm dụng SM. Hoạt động cao có nghĩa là các SM hầu như không bao giờ nhàn rỗi. Không bị thiếu dữ liệu (dataloader starvation) hay chờ đợi mạng, đó là phần thưởng cho việc tải xuống các shard cục bộ thay vì stream, điều vốn sẽ khiến chúng ta dễ bị tổn thương trước tình trạng treo mạng của hugging face. Mức chiếm dụng thấp là biểu hiện của các GEMM lớn chạy liên tiếp: các kernel matmul cố tình đánh đổi mức chiếm dụng để lấy kích thước register-tile lớn hơn. Bị giới hạn bởi tính toán và được cung cấp dữ liệu đầy đủ, đây là tín hiệu tuyệt vời cho thấy chúng ta đang sử dụng phần cứng tốt và có thể biến mỗi đô la chi tiêu thành một mô hình tốt hơn.

Đó là khoảng 1.047 TFLOP/s duy trì trên mỗi B200, hoặc ~25% MFU so với đỉnh FP8 dày đặc của Blackwell. (So với đỉnh bf16 thì là 50%, con số này quan trọng hơn một chút vì không phải tất cả các lớp tuyến tính đều chạy ở FP8.)

Chiến lược phân tán là DistributedDataParallel thông thường. Với 3,8B trên một node duy nhất, việc truyền gradient không bao giờ là điểm nghẽn, và cơ chế sharded-optimizer hóa ra là không cần thiết.

Tăng thông lượng

Thuê GPU không hề rẻ, khi làm việc bạn thường nghĩ về chất lượng mô hình trước chi phí. Khi đó là tiền túi của chính mình, thông lượng đột nhiên trở nên quan trọng hơn nhiều.

Điều này đòi hỏi công việc thực sự trên một chiếc RTX 5090 duy nhất, trước khi tôi thuê một node. Mô hình cơ sở 858M, bf16, đã biên dịch: 26.144 tok/s. Kết quả cuối cùng: 37.621 tok/s.

FP8 (+25%). Cả ba GEMM (1 forward và 2 backward) ở FP8 với dynamic tensorwise scaling. Yêu cầu SM90+ nhưng đó là một bước nhảy vọt về thông lượng khá tốt.

Đệm từ vựng (+33% tích lũy). Đệm 50.257 → 50.304 tốn 47 hàng embedding không sử dụng và mở khóa đường dẫn tensor-core nhanh. Gần như miễn phí.

Fused linear cross-entropy (+44% tích lũy). Liger’s FusedLinearCrossEntropyLoss hợp nhất matmul lm_head vào hàm loss và chia nhỏ bên trong, vì vậy tensor logits đầy đủ (B*T, vocab) không bao giờ được hiện thực hóa. Đo trực tiếp ở cùng batch size, nó chậm hơn 6%:

Mặc dù chậm hơn trên mỗi bước, nó lấy lại được một lượng VRAM đáng kể (8 GB trên 5090 của tôi) nên việc tăng micro-batch size bù đắp nhiều hơn cho 6% bị mất. Claude đã nhanh chóng từ chối nó vì nó thấp hơn 6%, nhưng nhìn chung đó là một cách tuyệt vời để giành lại thêm thông lượng.

MLP không cổng (Non-gated MLPs). Loại bỏ phép chiếu cổng (SwiGLU → relu², hai matmul thay vì ba) trên mô hình nhỏ: 183.035 → 214.173 tok/s và ít hơn 6 GB VRAM. Một lưu ý từ các thử nghiệm cắt bỏ (ablations): tỷ lệ trung gian SwiGLU là 2,75 không chuyển đổi sang relu² được. Mô hình học tệ hơn đáng kể. Hãy sử dụng 4× cho loại không cổng.

Trọng số chính bf16. Giữ trọng số chính của bộ tối ưu hóa ở bf16 thay vì fp32 giúp giảm 27% VRAM và tăng thông lượng từ 640K lên 1,4M tok/s trên cấu hình 1,5B. Đó là một sự tăng tốc khổng lồ, 2,2 lần. Chi phí chất lượng là có thật nhưng nhỏ: CORE 0,22 so với 0,23 ở 4.000 bước. Khi bạn tối ưu hóa khả năng trên mỗi đô la, việc xử lý dtype cẩn thận là một trong những đòn bẩy cao nhất và ít được thảo luận nhất.

Phần cứng. Cùng mã nguồn, mô hình 150M, FP8: RTX 5090 đạt 184.662 tok/s, B200 đạt 477.440 tok/s. 2,59 lần chỉ từ phần cứng, chưa tính đến việc VRAM dư thừa cho phép bạn đẩy batch size lên cao hơn nữa.

Những gì không hiệu quả

Masking ranh giới tài liệu với flex attention. Việc đóng gói các tài liệu vào một chuỗi cho phép các token chú ý (attend) qua các ranh giới, vì vậy tôi đã sửa nó một cách đúng đắn: ID tài liệu cho mỗi token và mask attention để mỗi token chỉ có thể chú ý đến tài liệu hiện tại của nó. Nó rất thanh lịch, nhưng tôi đã xóa sạch. Andrej Karpathy cũng phát hiện ra rằng sự rò rỉ giữa các tài liệu không làm mọi thứ tệ hơn nhiều dưới dạng đóng gói căn chỉnh theo BOS. Best-fit packing đã thay thế nó trong khoảng 10 dòng, và attention quay trở lại F.scaled_dot_product_attention(..., is_causal=True) không điều kiện. Tôi tin rằng điều này cũng phụ thuộc vào tập dữ liệu và các tài liệu huấn luyện.