Now what?

So I've finished watching Karpathy's "Zero to Hero" video series, now what? In his last video he built and trained a reproduction of GPT-2 on his GPUs. But my laptop isn't as powerful, and I could not train the model in a realistic time, So my next goal was to see what is actually possible on my Laptop.

Mork-1 (the plan)

Some years back I had scraped all messages from a private discord server me and my friends used. It totals to around 500k messages, from primarily three people. I was curious how well I could train an LLM with just this discord data

The Model

I used the GPT-2 reproduction I coded along to from the "Let's reproduce GPT-2 (124M)" video. First I changed some of the hyperparameters to decrease the size of the model:

	GPT-2	Mork-1
block_size	1024	512
n_layer	12	6
n_head	12	6
n_embd	768	384

Due to this I was able to increase the batch_size from 2 -> 32. And went from 12k tokens/sec -> 100k tokens/sec

The Dataset

I still had the old discord.json (~70Mb) file laying around somewhere. First it needed some parsing, I needed data in the simple format of.

Mart: Hey guys how are yall?
[Friend1]: Im good wbu?
[Friend2]: My dog died

Messages from bots needed to be excluded, mentions+channels+emojis are usually in the format like: <831233350593019914> So I simplified these to <url> <emoji> etc

The Tokenizer

I wrote and used a very simple BPE (Byte Pair Encoding) Tokenizer, But when trying to run it I found out its a bit slow and wouldve taken me ~8 hours to fully tokenize the dataset. So I asked AI to magically speed it up somehow. Also I decreased the vocab_size a lot. From ~50k in GPT-2, to just a vocab_size of 4096 for this project

Checkpointing

I implemented a simple checkpointing system, during the training run every 100 steps it saves the raw model to latest.py, and every 500 steps it creates a "step_000500.pt"-like file. This allows it to seamlessly continue training after stopping it mid-way or if it crashes. To successfully checkpoint the full state, I had to include the following data in the checkpoint:

• raw_model
• step, loss
• optimizer
• model_cfg
• train_cfg
• train_loader
• val_loader

Attempt-1

Training went fast and good! it completed the training run with a max_steps=1500 in just ~5 minutes! A fast learning start, that finally stabilises to a loss of around 3.6 You can see a recurring pattern in the train-loss, which is the training-data repeating. the max_steps=1500 has an epoch of I think ~7

Generation

But now, what is most important, what text can it generate! I wrote a simple generation.py CLI tool, and these were some of the first results:

--- sample 2 ---
Marto_0: O
Marto_0: Wait which
Marto_0: Im not that tough I just to go tough a mth 
Marto_0: Yeah 
Marto_0: Its a okay 
Marto_0: Which does you dont well not really knows more tough 
Marto_0: And its just not really 
Marto_0: Idk 
Marto_0: They are a ting an hind me me tough it just becond to give me of a higher 
Marto_0: Ooo long is just the whole pety capple croid and you were you could 

Sooo... could be better haha, but if we analyse what it already does correctly: It mostly uses english words that do exist, it creates somewhat valid sentences. Considering the LLM had to learn the english language from scratch purely through discord messages, I think it does pretty decent.

Attempt-2

AI told me I won't be able to get any better results... Told me the issue is the small dataset, but I didn't believe so! Since I was able to complete a full training-run in just 5 minutes, I could easily increase the model size again:

	Mork-1	Mork-1.1
n_layer	6	8
n_head	6	8
n_embd	384	512

Which decreased the tokens/sec from 100k to 60k. I also changed the max_steps from 1500 to 3000 for a longer training-run.

Training Shuffling

In the previous loss-graph you could see the same dataset re-appear. Not sure if it matters greatly, but I split the training up into chunks, which are shuffled every epoch. This should stabilize training a bit more.

Results:

WOW, this blew my expectations:

Notice I also plotted the learning-rate in green, the reason it spikes at 1500 is because thats when I changed the max_steps to let it run longer. Altough accidently, apparantly this is called a "sawtooth learning rate", which can help neural nets escape local minima. However in my case it doesnt seem to have a beneficial influence on the loss. The best-val went from 3.6 from the previous run, to 3.2!

Here are some new samples generated from this improved model:

Marto_0: oh god
Marto_0: Lmao
Marto_0: I think it is
Marto_0: But yeah
Marto_0: And yeah I do be wondering
Marto_0: I mean
SuperVK: idk
SuperVK: yea I guess I think
SuperVK: but idk
Marto_0: Lmao
Marto_0: I think its quite good toorne
Marto_0: Its just more of copy on
Marto_0: Like
SuperVK: yea
SuperVK: but not that I wanna play
Marto_0: fair point
Marto_0: I mean, how faile

Can't say it generated a highly intellectual conversation, or anything comprehensive even haha. But I think it does quite accurately reflect how I used to chat at that time (I was aged 14-17).

Attempt-3

I realised the dataset was still quite dirty, and needed some cleaning. There were many many issues.

Val / Train mismatch

the val split was just the last 5% of messages by Time, but our chat behavior changed by time aswell. This caused val-messages to be 27% longer on average compared to train-messages. The composition of users in either split also changed, with one user going from 24% of train-data, to just 5% of the val-data. This makes it so that the training can never accurately predict the val

Bots and URLs not fully filtered

There were still a few bots left unfiltered that cluttered the training set with repetitive sentences. Many things like IDs were filtered out, but URLs werent (images, GIFs etc). These URLs are unpredictable random strings, and thus hurt the training. And apparantly 8% of final tokens were due to these URLs!!

Pre-tokenization

One improvement thats often applied, is splitting the text before tokenizing. Usually splitting by space, special characters etc. Otherwise weird issues can occur like multiple words forming a single token, maybe "that is" becomes a single token, which degrades understanding for the LLM afaik.

Reserved tokens

Some words come back a lot in the dataset, for example my username "Marto_0". This then causes the tokenizer to make multiple merges like "Ma"/"Mar"/"Mart"/"Marto"/"Marto_"/"Marto_0" etc. A simple fix is to just reserve tokens for the top-10 usernames, and have "Marto_0: " be a single token without filling up merges.

Dataset changes:

I made various changes to the tokenisation and dataset parsing, which caused different results in the tokenized data. The total amount of Train tokens decreased due to filtering out more bad data. And the Chars/token increased, which means a token is now around ~20% longer. The positive is that the AI model can now have a larger context window. One annoying thing tho is that the validation-loss and training-loss can no longer be accurately compared to previous runs on the old tokenizer.

Val/test split fix:

I fixed the Train / Val split mismatch. Instead of using the latest data as validation, the validation data is now constructed by taking random time chunks out of the training data.

Now each user is represented at roughly the same percentage in the training split and the validation split

Old --> New

Results:

WOW, first noticable change is that the variation in train-loss has decreased quite a lot. Likely due to the cleaner dataset (no URL prediction etc).

Secondly, the trajectory of the val-loss is a lot cleaner, which I think means its more aligned to what is being learned in the training set.

Thirdly, the best val-loss is actually ~3.5, higher then our previous 3.2, Sadly this isn't easily comparable anymore because of the higher Chars/token, because one token holds more information now, its also harder to predict that token. I'm pretty sure the model did actually improve.

Fourthly, There is a very visible widening gap between the validation loss and the training loss. That means that the model has started overfitting on the training-set. Basically it started learning exactly what is inside its training, and doesnt actually get better in predicting outside of that data.

I think this is a clear indicator that the dataset is too small now, and due to that theres not much left to improve without more data

Samples:

The samples have gotten somewhat more comprehensive, theres even some talk about the simulation theory lol:

Marto_0: But they can't just make it work
Xeukxz: But the universe cant be the best
Marto_0: And people want to make it more of the simulation. 
 Just a simulation that learn the simulation. And it can change the universe
Marto_0: And if you can learn that.
Xeukxz: Well the universe is the universe who make it
Marto_0: But if you want to learn it. If its a universe. 
 You can just do it. But its not that particularity that it can work
Marto_0: And make it can be any reason the universe can make its own universe

Or some yapping about coding:

Marto_0: but yeah okay
Xeukxz: :emoji:
Marto_0: so how would you want to make a code that runs on an example
SuperVK: idk then
SuperVK: but maybe it just doesn't work?
Marto_0: yeah
Marto_0: I should prob just do the whole code to work with the console.log or something
SuperVK: cuz then u can make a new project to get into the code
SuperVK: or like that
Marto_0: yeah okay
SuperVK: or just try it in the code
Marto_0: idk
Marto_0: I'm just gonna work on it now

The sentences don't always make sense, and the full conversation makes even less sense. But still, I'm surprised how well it actually does generate english. It can keep topics constant, like coding, or simulations. And the responses occasionally make sense.

Dropout

Actually, before being fully satisfied with the current results, I still wanted to try one more trick.

Dropout is a technique where you turn off a certain percentage of neurons during training. This somehow forces the LLM to actually learn patterns instead of learning and overfitting on the data. I ran the full training again, with a dropout of 10%. And surprisingly this really worked!

The best val-loss went from 3.5 to 3.35! In the graph you can see that the train-loss and val-loss are pulled closer togheter. Meaning, less overfitting is happening. The val-loss actually isnt that far away from the train-loss in the end, maybe increasing dropout to 15% can close this gap even further, but there doesnt seem to be much to gain.

Attempt-5

So I've really squeezed out all performance I can from the limited personal dataset. There is still a way to get it to be better tho, to pretrain it on a way larger dataset, then finetune it on my own dataset. This way it can actually learn proper english and understand the language, and afterwards train it to learn our way of speaking.

This would work the best if the pretrain dataset is as similar as possible to my dataset. I decided to use this dataset Discord-Dialogues on huggingface.

I then need to split the training-run into two parts, the pretraining, and the finetuning. For now I just do it simple sequentially.

I also upped the max_steps to 4000, I think especially the pretraining needs this since the dataset is so much bigger.

Results:

Pretrain:

Finetune:

Not bad! final val_loss after finetuning was 3.01! down from 3.35 without pretraining. And just the pretraining model reached a val_loss of 3.3 (on its own dataset only, not my own dataset)

Its interesting to see how much quicker the model learns having done pretraining, It already reaches near best-val at step 1000, while without pretraining it used to still improve at step 3000.

You can also see that unlike the finetuning stage, where you can see overfitting slowly happening (train-loss getting lower then val-loss). However in the pretraining stage both loss lines stay near, and no overfitting is happening, which makes sense since the dataset is way bigger and theres not much data repetition

Now some more samples! Firstly, samples generated from the best_val pretrained model (but without any finetuning)

Pretrain samples

A: I need to learn to learn how to play
B: Just do it to a game
A: I mean, I'm not able to

A: i think i should get an extra 1 for this team
B: After your 2 pulls and you'll get the 2nd one, but I would get a 1k.
A: i mean i'm gonna do it for the 2nd time

A: Is there a difference between the two that adds more orbs?
B: The only difference is the two
A: So it's a good option?

A: You can get to the rules in the server
B: Yeah but you can still get to the point where it is
A: See what, I can't find the point where the rules are going to be, 
 but I'm not sure if that's a problem

A: Hiii
B: Haiiiiiiii
A: Hru
B: I’m good just got home and sleepy, and I’m doing pretty good

Looks like pretty generic discord conversations! not bad, its clearly learned english pretty decently already.

Samples after finetuning:

Marto_0: Whatsup guys
Marto_0: Oh shitt
SuperVK: yo @user u wanna play rainbow?
Marto_0: no
Marto_0: just had 2 hours of sleep
SuperVK: what?
SuperVK: oooooh
SuperVK: what was it?
Marto_0: a little bit
SuperVK: what are you gonna do?
Xeukxz: School
Marto_0: oh
Marto_0: wait nvm
SuperVK: oh
SuperVK: what's the last thing u played?
Marto_0: a little bit late
Marto_0: But I got like 6 hours of sleep
Marto_0: And I had to do a 2 hour walk so now I can sleep at like 8pm
Marto_0: But I had 3 hours of sleep
SuperVK: oooooh that's nice
SuperVK: so what's your next job?

Sounds kinda accurate haha.

BUT, im not done...

Rotary Positional Embedding (RoPE)

Instead of the previously used "positional embedding", I implemented RoPE instead. Its a newer technique for embedding positional data into the LLM. I wont explain it in detail, but its great because it sortof remains relative positional data between tokens, instead of absolute positions.

Anyway I wanted to see how this improves the result, so this is the only thing I changed from last run. I also just did the pretraining run, not the finetuning run yet.

Very nice! learns faster and converges to a lower val_loss! We improved the val_loss from ~3.3 to 3.20, nothing crazy but better then nothing!

And...

I also replaced LayerNorm with RMSNorm. And I used SwiGLU instead of GELU. However these didn't produce notable improvements.

Auto-Experiments

Pretrain runs were taking around ~20 minutes on average, which is a pretty nice and short time. It allows me to run many experiments to find optimal hyperparameters. I set up a system where I can make json files that contain certain experiments, for example:

{
  "id": "depth_screen_1k",
  "base": {
    "max_steps": 1000,
    "data_dir": "data/dialogues"
  },
  "sweep": {
    "model.n_layer": [8, 10, 12]
  }
}

This would do 3 seperate pre-training runs, until 1k steps. The whole setting is the same besides that n_layer will be changed. The hope is that with 3 data points I get a general idea of how this hyperparameter changes the learning.

I also created a queue.json file, where I can list multiple experiment files. I can then run the queue, and the script is able to run for hours, doing multiple runs for many different experiments.

This way im not forced to constantly pay attention, change configs, compare results, restart runs. Instead I can just queue experiments overnight, let it run for hours, and come back to the results!

For the experiment I just showed, these were the results:

You can see that the validation loss got slightly better with more hidden layers (n_layer), but honestly the improvement is barely visible. In the second screenshot you can see that L12 took almost 45% more time to train compared to L8, in this case its clearly not worth it. I should run another experiment to test with 6 hidden layers, see if this still gets to the same result but with faster training.

Different Learning rates

Another hyperparameter I wanted to experiment with was the learning-rate. And these showed big results!

Very different end-results, and no difference in training-time! This seems like a change very worthy of making However I dont want to immediatly make conclusions and change hyperparameters too soon. Remember this just shows a training of 1k steps, while to max best_val we usually need 4k steps of training.

Also we now know the highest LR of 1e-3 has the best results, but what about an even higher one? We should experiment with multiple higher LRs, to see at what point it starts to degrade training. This way we can find the optimum.

Anyway, I ran over 6 experiments, testing all sorts of changes to hyperparameters. Honestly, not as much result as I had hoped. I tried scaling up the model to bigger sizes, but this only came with a tiny improvement, but took 50% longer to train.

Two changes sticked tho, The learning-rate could be a lot higher, I found the best is at 2e-3. And other change, which I couldve easily predicted honestly, is to turn off dropout during pretraining, since theres no problem of overfitting with the large pretraining dataset. After a longer run of 6k steps, I went from 3.20 to 3.07 val. I'll take it.