This project demonstrates the feasibility of training a high-performance 3.848B-parameter decoder-only language model for under $1,000. The final model, "little-lm," achieved a 0.384 score on the CORE benchmark after being trained on 65.3B tokens over 43 hours using eight rented NVIDIA B200 GPUs. The architecture utilizes a Llama-style configuration including RMSNorm, RoPE, Grouped Query Attention (24 query/8 KV heads), and relu² MLPs, supplemented by 721.2M parameters of value embeddings (14 tables) which represent 19% of the total parameter count. Key technical optimizations included a trapezoidal learning rate schedule with a 50% linear cooldown, the use of the Muon optimizer for matrix parameters combined with AdamW for others, and FP8 training via torch._scaled_mm to increase throughput by 33%. Additionally, the author implemented vocab padding to a multiple of 64 and utilized Liger’s FusedLinearCrossEntropyLoss to reduce VRAM consumption, allowing for larger micro-batch sizes. The study highlights that increasing context length from 1024 to 2048 tokens significantly improved CORE scores, particularly for context-dependent tasks like SQuAD and BoolQ, despite a 9% reduction in training throughput.
Read original
hackernews