LiteGPT: Pre-training a 124M SLM on a laptop (without melting it)

TL;DR: I pre-trained a 124M GPT class model from scratch using a consumer-grade RTX 4090 laptop by processing 10 billion tokens locally, implementing custom thermal profiles, enabling GPU Overclocking, optimizing Model Flops Utilization (MFU) to ~15%, and implementing a full SFT pipeline (Alpaca) to turn random noise into a functional Instruction-Following SLM.

Why this project you ask?

So I was tinkering with my recent research work which involved anomaly detection using Transformers for which I was utilizing a couple of open-source LLMs like Llama, Mistral and Gemma. But then I thought, how difficult would it be to build an LLM of my own because my research model didn’t require a powerful pre-trained Transformer. Instead of reusing, can’t I build something of my own that will give me more engineering flexibility? I don’t have access to any of the powerful GPUs and my laptop is not a cutting-edge PC with a GPU stack. So, obviously cloud right? Nah, let’s experiment with what I have now. But first, I need to know whether this might work or not. To make that clear, let's try building a simple LLM first or SLM (Small Language Model).

But I had no idea where to start because I read a lot about deep learning algorithms and implemented a lot of them but never have I explored how to build a transformer. After hours of research, I came across Andrej Karpathy’s nanoGPT. It’s old but was more than enough for my requirement. I also knew that he started to rewrite his codebase to nanoChat, his new work. But nanoGPT had more documentation and reddit posts that I can read from, sooo yeahhhh. I took this route.

Constraints: Stoppers I was facing initially

I wanted to build a POC on my own laptop. Let me list down the system specifications.

The Architecture: How I wanted the output to be?

A GPT based Transformer (or simply an LLM) of 124M parameters because it’s the default provided by Andrej’s examples and I am building just a POC. 12 Layers, 12 Heads, 768 Embedding dimensions.

So, I wanted to utilize the same. Recently, I came across a paper where I discovered the FineWeb Dataset which had all the quality blogs scraped and stored as a knowledge dataset. But that’s a really big dataset with about 18.5 Trillion tokens. I don’t have the patience to build one of that size. So, I chose FineWeb’s sample-10B subset.

The tokenizer is GPT-2 Byte-Pair Encoding (BPE). This is provided by default so, I didn’t want to delve much deeper into this territory.

At the end, once this entirely runs, we would get a basic text completion engine that uses GPT architecture.

Engineering: Here we go…

To simply put, the primary engineering effort went to configuring the hyperparameters of the model itself.

out_dir = 'out-fineweb'
eval_interval = 500
log_interval = 10
eval_iters = 200
wandb_log = True
wandb_project = 'fineweb-gpt2'
wandb_run_name = 'gpt2-124M-10B'

dataset = 'fineweb'
init_from = 'resume'

batch_size = 10
block_size = 1024
gradient_accumulation_steps = 48

max_iters = 20000
lr_decay_iters = 20000

n_layer = 12
n_head = 12
n_embd = 768
dropout = 0.0

learning_rate = 6e-4
min_lr = 6e-5
warmup_iters = 2000

device = 'cuda'
compile = False
dtype = 'bfloat16'

This is the primary configuration file config/train_fineweb.py. These are all the default configurations given for the openwebtext dataset. Not a major change. Changed the dtype from float32 to bfloat16 because of the out-of-memory exception that occurs due to the lack of memory. Increased learning rate because of the vast dataset that was presented to the model, just doubled the number that already existed. But all these aside, coding / configuration was the least problematic issue in this process because training the model gave me more issues than trying to build a house in Bangalore.

Issue 1: Data Preparation Constraints

Preparing the dataset was taking a lot more time than expected due to single-threaded implementation. So, I implemented parallel processing to speed this up. This prepared the dataset in about a few minutes.

Issue 2: Shared Memory

The first execution gave me an unexpected problem. It would take more than about 6 days for my model to complete. I mean, I expected a long time for training this model but not this long. So, I started to look into the task manager and the model took about 17.2 GB memory cause in Windows there is one of the best features ever implemented called Shared GPU memory which means, the RAM will also be taken as a GPU memory if there is not enough video memory available. When the shared memory is utilized, the compute speed of the GPU will drastically reduce. So, I had to reduce the batch_size to 10 to fit into the GPU memory itself. This brought the VRAM requirement of the pre-training process to 15.2 GB.

Issue 3: Thermal Optimzation Requirements

Initially, the system said it might take more than 48 hours for the pre-training to complete. But it was constantly wavering to 54 hours and even more. When I checked the nvidia-smi option, I observed that my GPU had started to throttle because the thermals were reaching more than 87 Celsius, which is the upper limit for most of the Mobile GPUs. So I had to modify the voltage supplied to the GPU using the nvidia-smi tool. There is no documented way of identifying the optimal power consumption. So, it was just trial and error.

sudo nvidia-smi -q -d POWER

Used this command to know about the default power usage. After constantly running the model in various custom profiles for Power Limits, I found that 150w of supply + 15w for Shared Supply to CPU and GPU was the optimum temperature which yielded about 15.81 MFU and 8.49 ms per iteration and I could see the model getting completed in 40 hours. Also, I had to do so much more like cranking up the fan speed of the system to the max and laptop stand for better ventilation. This was an interesting problem to solve.

Issue 4: Power Cuts

Well, so basically I wanted the model to run for 40 hours consistently but I was pretty sure that there would be power cuts interrupting the model. So, I had to implement a checkpoint save mechanism that will save every 500 iterations. About 20,000 iterations were being run for a successful execution of this model. This is one of the much needed features for my situation.

The "Brain Transplant": Instruction Tuning (SFT)

So, the pre-training was done and proudly successful. It gave me an average result of ~15% MFU, 8.54 ms of iteration time, completed 20,000 iterations in 40 hours with a training and validation loss of 3.04 and 3.07 respectively. So, I was assured that this model was not overfitted with the dataset. I was able to see other’s results on the same implementation that had only about 12 MFU on average. But through sheer will, determination and through perseverance... Nothing of that sort. It’s just the optimization in play did its job perfectly. But then, what we built was just a base model. It’s just a Next-Token Predictor model also known as Auto-Regressive Transformer that will predict the next occurring word through statistical probability. But I wanted the model to respond to me, not complete what I am saying. So, lets work on this again.

To migrate the model from “What could be the next word?” to “Answer me goddamnit”, we are going to use the famous approach of Supervised Fine-Tuning. Our model has knowledge from the training. Now we are going to make the model respond to us based on a structure. So, we are going to choose the Alpaca format, why do you ask? Because it was the first google search result. Just kidding. It had more documentation and references. This dataset has a specific format it follows. For example:

Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{instruction}

### Input:
{input}

### Response:
{output}<|endoftext|>

Through this process, I would be able to converse with my SLM like my dyslexic buddy over a beer.

Results & "The Hallucination Reality"

First, we'll look into the metrics of the pre-training model. So, as claimed previously, this model was pre-trained on a consumer grade mobile GPU and had a limited memory of 16 GB for the training process. So, I configure the batch_size to be contained within this memory limit through trial and error process starting from 32. Finally, I found the optimal configuration to be 10.

This is the training loss of the pre-training model. This represents how much the model is learning from the dataset. The lower the better.

Now, for the sample output of the auto-regressive model. This model is just a prediction engine, which identifies the next probable word based on the previous words. Let's see what output it gives.

As you can see, the model is not trying to fact check the input or communicate with me or enquire about my marital life like a distant relative you don't even know (It's an Indian joke). It's simply pushing out the best words it knows, that it predicted to be the best.

Now, let's see the sample output of the instruction tuned model.

Now, this instruct model is fine-tuned with knowledge and action based responses that makes the model to perform an action based on the input instruction. This model is more useful and more relevant to current AI models.

Conclusion: What I Learned

• This model is still dumb. How or why you ask? This doesn't have enough knowledge base to even compete with the other basic models such as BERT which is a 437M parameter model. There is no defined number of tokens that can make a model a general purpose model. It's basically trial and error.
• It hallucinates too much when unknown subjects are part of the input. Again, the reason being the model is still dumb. It gaslights me so hard I started questioning if I was the one who didn't know the capital of France.
• Reaching 15% MFU on a consumer card is a badge of honor. It proves that the bottleneck for democratizing AI isn't math—it's heat dissipation and VRAM. If you have a gaming laptop, you have a research lab. You just need to be willing to treat your GPU fans like jet engines.
• 30% of the model memory is purely dictionary of words that the model has learned and 70% is the attention model knowledge probabilities.

The Numbers: What 124M Parameters Actually Looks Like

According to my transformer_sizing analysis, here is where my laptop's memory went:

• Total Parameters: 124,337,664
• The "Dictionary" (Embeddings): ~38M params (31% of the model!)
• The "Brain" (Attention + MLP): ~85M params
• Final Model Size: 0.5 GB on disk