My Learnings From Training LLM from Scratch

Research on training LLM from scratch is fun, but at the same time it’s frustrating because there’s lot of factors contribute which is difficult to find which factor making the model marginal better and less, And that’s the discovery in whole Deep Learning! If you read any pre-trained LLM architecture, that will be around ~200-300 lines code, but if you go to understand that particular lines, a 300 page book will be less.

in any model training, these main things mainly contribute:

1. Architecture of the Model
2. Data
3. Optimizer (controls the learning)
4. Learning algorithm (must be gradient-decent till now)

The fastest way to learn “Training LLM” is the slowest way, which is building it from scratch, tinkering it, diagnosing it why it works and why not. There is no way where you can read and watch a video and can next day start training the SOTA LLM. But while tinkering you can fail fast and learn fast.

I trained an ~160M Model combining different architectures like MTP (multi-token-prediction), HRM (Hierarchical Reasoning Model) and causal LM like GPT. And it even did better than nanoGPT and surpassed the CORE accuracy.

After doing 100s (I don’t remember, maybe ~300), Many didn’t work, and the ones which gotten best results I took them and built over them. I read about many architectures like text-diffusion, HRM, MOE, etc and honestly I tried all to build a ~150M version of it, but some architecture didn’t work like diffusion, that require long training and has slow convergence. Another thing I wanted to train quickly with less tokens (like ~20-30min), similiar to nanochat speedrun, as there are small models which are trained over 40B~50B tokens but I wanted to just train fast and small in around ~2B-3B tokens.

Learnings

First thing I learned that if I want to experiment an idea, how can I simplify it and test it fast, the mistake I did earlier that I tried replicating full pipeline, If I am taking an inspiration idea from from a paper like HRM, I found the core mechanics, which is the “cycle loop”, that gives latent reasoning, and asking question “How can I take this core idea and put it and test it, instead of making a larger change in the architecture?”. I understood why HRM works, and took the inspiration to build a simple loop in my architecture and test it. And it worked, it gave me some extra CORE score, without increasing too many parameters and major architecture change.

Second thing is that one thing is clear that small llm models like ~160M can become good in narrow task and not in multi-task. I tried doing SFT on math dataset and they learned the math like gsm8k but wasn’t gone bad on other task. It is because they have not enough neurons to capture those information, as we know LLM is just a compression algorithm, which compress all the knowledge into it’s latent space, but that latent space should be big enough to capture those information, that’s why scaling laws says if you scale the transformer model (add more parameters), it will eventually eat all the knowledge and become expert in it. Well this compression comes down to architecture as well where the better architecture compresses the knowledge better and other doesn’t that’s why I explore a lot in the architecture space and assembled different inspirational ideas which made the model better.

Third learning was that RL only works when the model has more than above general knowledge of that particular topic, otherwise it won’t going to benefit there. I wanted to post-trained the model after SFT on my model with the goal that the model will learn the gsm8k dataset, but unfortunately the SFT model wasn’t good at math at all, it didn’t learn, it collected some rewards on some math problems but couldn’t solve most of it, training the forever could make the model eventually figure out at higher temperature but that’s not what we want, that’s why we pre-train or do SFT on the model more on math data so it could have better knowledge and reasoning like an average, then RL could make the model expert (a cherry on the top). In-short, RL can’t make someone who has little to no knowledge about a field, but can make expert if it is more than average.

The intiution is that RL gives reward 0 to wrong and 1 to correct, you can also add intermediate steps like steps and formatting but generally in simple way it’s very sparse so if it doesn’t have a internal reasoning and knowledge to how to solve that problem, it will only rely on trial and error, hitting and getting rewards, and that will led the model to search the entire token space to predict those combination of tokens which will lead to rewards, so it’s not reliable.

Fourth is, the four building blocks I defined in starting: architecture, data, optimizer and learning algorithm isn’t only thing that help the model train better, but also Infra. A good infra is needed to train and scale a model which for either doing early experiments or production level training. When I was training, I wasted a lot of gpu minutes because of non optimal infra code which is not efficient which led to do long minutes training, waiting time, cold start, issues in gradient sync in distributed training. Systems knowledge is very important for doing the research in LLM space.

Important thing to make sure to utilize the full potential of the GPU, I was doing RL experiments, and saw, training only utilizing the 38% of the GPU, because in RL you collect rollouts and then you do backward pass on that rollouts, where rollouts are getting collected sequentially for each example, and it was not optimal infra for a GPU, so I had to stack all the generation in a single batch where all the generation going to happen next token prediction but 256 samples are predicting the next token at the same time, that lead to generate synthetic data fast for training.

Fifth and the last but not the least is that documenting the experiments, and be logical rather than emotional, take decision based on benchmarks. initally whenever I had gotten idea, and curious to try it, I just make different experiments of that idea but I do not document it rather just be emotional “ok let’s try this if it can work doing this, let’s try increasning this maybe it can get better than this”. And I wasted so much time doing like this, it is not that I didn’t figure out a way to do proper research but it is just I was too curious that I skipped to follow the proper process, but sooner I understood and created a leaderboard, and started ranking and documenting experiments.

That’s for all today, thanks for reading this! find the code here: https://github.com/harshbhatt7585/deep-learning-papers-implementation/blob/main/tinyGroot