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 4.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-messages, to just 5%. 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 thee 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.