Mork-2, Learning Python

It's been a while

My last blog regarding AI was making mork1, which was about 2 and a half months ago. I've been a bit quiet since, mostly due to moving to study abroad in China.

I had actually started a few new AI projects, with the goal of learning SFT, LoRA, and other post-training methods, using a small base Qwen model. However I let AI do too much on its own, and I had gotten lost in my own projects a few times.

One issue with post-training an existing Qwen base model is that it's already very capable. Even the instruct/base model already writes good Python, follows user/assistant instructions, etcetc. So there are not many easy tasks to train on.

So that's why I decided to make mork2 instead, a new model pretrained from scratch. This way I still understand the whole stack, and I control exactly what the model has and hasn't learned. I think this will be a better learning experience.

The code is public: Mart100/mork-2.

Mork2 plan

Mork1 was just about learning English, and replicating Discord messages written by me and my friends. I want to take mork2 a bit more seriously, though. I want to train it on a lot more English, so it has more world knowledge, and in the end I hope to train it on a lot of Python code. Eventually I also want to post-train it to get even better at Python.

The model

The model architecture has mostly stayed the same as in mork1, but I increased its size:

	Mork-1.1	Mork-2
n_layer	8	12
n_head	8	12
n_embd	512	768
vocab	4096	16384
total_params	29M	98M

The tokens/sec while training has decreased from 60k to 14k, so this model will need a lot longer training runs.

Data

I primarily want this model to be able to write Python functions, but for that it also needs to follow English instructions. Therefore it needs to learn both English text and Python code. I decided to go for the same composition of data as SmolLM-Corpus:

• FineWeb-Edu: a dataset of English text that mostly contains high-quality educational webpages. This primarily tries to teach the model a large range of facts and knowledge.
• Cosmopedia v2: this dataset mostly contains stories, blog posts, and textbooks. It mostly teaches use of the English language.
• Python-Edu: just a bunch of high-quality Python code.

Together these should be able to train a model that has general world knowledge, English knowledge, and Python code knowledge.

First Training

The whole training loop has generally been kept the same as in Mork-1, so after the model size changes and different training data, it's ready to start its first run:

Not bad! In just 4k steps, which took about 2.5 hours, the validation loss fell from 6.31 -> 2.71.

This was just 131 million tokens, though, which comes out to 1.3 tokens per parameter. That's still very low. Usually 20+ tokens per parameter is recommended.

Training on the rest of the data (8GB in total) would take about 60k steps, around 40 hours on my laptop.

One issue here tho: the orange line, the validation loss, is very noisy, so it's hard to judge if the model is consistently improving.

This is because after every 100 steps, a different slice of the validation data is used, and the slice is quite small: 102,400 tokens. Sometimes the slice holds a lot more difficult-to-predict tokens than another time.

An easy fix is to increase the slice size to 512,000 tokens. The only downside is that it takes around 30 seconds to calculate every 100 steps.

Overnight Training

I first let it train for 2x2 more hours, and then a long run of 8 hours overnight:

Wow, we got to 18k steps in total.

Firstly you might notice the green line (the learning rate). It accidentally reset to its starting learning rate, instead of continuing where it left off.

It's clear that this high learning rate was bad. Both the training loss and validation loss shot up during this time. I reset it twice before I was happy with the learning rate.

Also you can see the orange line is a lot more stable after step 6k, due to the changes I mentioned earlier. Later on I managed to smooth it even further by using a fixed val instead of a rolling val.

Now the line from 6k to 18k might look quite flat to you, but it is in fact consistently slightly improving. We just need to zoom in a bit to see it...

It shows a steady downwards line from a loss of ~2.7 to ~2.5.

But you'll probably wonder, what does that even mean? Magic number goes down a bit, wow congratz. And well, I have no clue. It just means the LLM got slightly better at predicting the next token for the data.

It's time, though, to put the model to the test and see how well it can actually generate English and Python now. I wrote a very basic eval (evaluation) that asks the model to write 20 very simple Python functions. Here's an example:

Prompt:

def reverse_string(s):
    """Return the string s reversed.
    >>> reverse_string("hello")
    'olleh'
    """

The model then has to complete the function. Whatever function it generates gets loaded into Python, and then gets tested with two examples to see whether it works:

welcome > emoclew
today > yadot

This way the model can write the function however it wants. As long as it works, it passes the test.

I tested the model on 20 of these examples, and the model tries each example 10 times. Out of the 200 total functions, 23.5% passed! And 13 of 20 tasks passed at least once!

This gives us a way better idea of what the model can do. It can write some of the most basic Python functions. Not a very smart model yet, since these functions might even have been in the training set, but it's a decent start.

This gives us an idea of where the model's capabilities are currently at, but what is even better is to see previous scores as well, so that we can see the trend. Luckily a snapshot of the model gets saved every 1k steps. This means we can still use the models at these steps to calculate its performance.

And one more thing. Currently the data the model trains on is a combination of 3 datasets, and we graph the loss of each of these. But it would also be useful to see the loss of each dataset individually.

The first graph shows us the loss during the entire training of each dataset individually. The loss of each dataset is quite far apart, but this makes sense.

Python code is actually easiest to predict, because a lot of the syntax, indentation, and common idioms repeat. FineWeb-Edu has the highest loss, and is the most difficult to predict, because it contains a lot of difficult knowledge. Cosmopedia is synthetic text written by another LLM, so it's very regular.

The trajectory of each is actually very similar. This makes sense because each dataset has an almost equal split during training.

The second graph shows us the Python functions eval I explained earlier. Most important are the blue and green lines.

The blue line means: if each test is attempted 10 times, how many tests pass at least once. The green line means: out of all attempts, how many passed. These lines are both steadily increasing, which is great.

Remember how the loss went from 6.3 -> 2.7 during the first 4k steps, and only 2.7 -> 2.5 during the 14k steps thereafter. Even tho it seems like the model barely improved during the second run, when we look at this graph we can see that it got twice as good at writing Python functions during that second part.

That tells us loss doesn't tell the full story, and evaluation gives you a better idea of its capabilities.

Cooling down

I used AI to analyze the current state of the model, and asked it to recommend what to try next. It told me I can try "cooling down" the model. I had to learn what this meant.

Basically, at the end of training you can cool the model down. It means you do a last stretch of training, where you let the learning rate decrease to 0. This means the weights will stop jumping around a lot, and they will instead settle down. It won't learn much new, but it will stabilize its current knowledge. A cooled version of a model is a finished one. It's the last step of pre-training.

Here are my results using this technique:

The loss decreased substantially during the cooldown that took 1k steps. It decreased a lot more than it did during the training before the cooldown.

On the Python eval (passed attempts out of total attempts), it also went from 23.5% to 31.5%. This is a very welcome improvement, considering the cooldown run took just 30 minutes.

But now... it's time... for more training!

Another night to 30k steps

After ~8 more hours of overnight training, we have reached 30k steps. Including a cooldown to 32k steps.

Results:

The mixed loss went down from 2.4 (cool_18k) -> 2.3 (cool_30k). It's better than nothing, but the gains are slowly fading now, and I don't think there's much reason to continue training more steps.

At some point the model will reach its limit, since it's a very small model (98M). For comparison, modern LLMs have multiple trillion params, that's 10,000 times as much!!

The model is still consistently improving during training, but the gains become smaller and smaller. And training longer on my laptop GPU especially isn't worth it anymore.

What about the evals?

Well... those are not going great, from 31.5% to 35%. The eval only consists of 20 tests tho, so it's not very accurate in measuring performance.

I tried to eval on the public MBPP (Mostly Basic Python Problems Dataset), which consists of 974 tests, of which we test 257 tasks. I also made a second custom eval with 160 tests with mid-difficulty problems.

They slightly improved between cool_18k and cool_30k, but not a lot.

• MBPP: 1.25% -> 1.6%
• Custom eval: 1% -> 1.4%

(sampled pass@1)

Both of these new evals are still too difficult for the model, so we can't accurately say it improved.

SFT TIME!

Most parts of this project so far have been similar to mork1, refining the pre-training step. But now it's time to apply some SFT post-training.

SFT stands for Supervised Fine-Tuning. It usually means fine-tuning a pretrained model on a more specific dataset, or training it to follow a specific format or instructions.

Until now the model has just learned to complete any English text or Python code. But it hasn't been told that, if you ask it to write a specific Python function, it should actually follow the instructions. It's only just learned to write Python, not to follow instructions.

We could potentially massively improve the model's scores on the benchmarks by simply teaching it to follow instructions better. Besides that, my plan is to do SFT exactly the same way as pre-training. The only difference is a mask on the prompt. This means the loss is only calculated on the reply, and not the prompt.

The data

We again use a split of three separate datasets:

• opc-sft-stage2 / educational_instruct: the most important of the three. It contains English instructions describing a Python function, and the Python code that implements it.
• sc2-exec-filter: a more difficult and more diverse dataset of instructions + Python code.
• smol-smoltalk: pure English text, not code. It contains a long list of messages in a user-assistant style.

These combined should allow the model to effectively learn to follow user instructions. Especially Python function instructions.

We do have to filter the data somewhat. The biggest issue currently is eval contamination.

The educational_instruct dataset contains the same functions as the MBPP eval. Once the model learns those exact functions, it will score very well on MBPP. But then it stops telling us how good the model is overall. It already knows the answers to the exam.

We should also filter duplicated data, data that doesn't fit inside the context window, Python code that uses third-party imports, and SmolTalk data that doesn't include a single user-assistant turn.

After that, we're left with about 27M tokens of data. The model will train for 2 epochs (train twice on the whole dataset). This should take approximately 2 hours.

The run

Looks very good, I think.

A few things to note. There's an obvious dip in the training loss in the middle. This is when the model starts the second epoch. It starts overfitting on these specific examples, but since the val loss is still decreasing as well, this is no problem.

Secondly, you can see the cooldown at the end again, and how it helps the val loss drop even faster.

And thirdly, the loss is actually really low, 0.7 at the end. Note that this is not comparable to the pre-training loss. This loss is calculated on just the reply, not the prompt.

The results

Wowie!

Smoke test from 60% -> 90%, mid eval from 3% -> 39%, and MBPP from 4% -> 19%!!

That is a way bigger improvement than I was expecting.

For the mid eval we tested both "instruction" and "function stub". "instruction" is just asking the model to write a specific Python function with a specific name. A "function stub" already includes the start of the Python function, and asks the model to complete it.

Smoke and mid were originally both function-stub style, but educational_instruct was fully instruction style. It's good to see that even tho the model was not further trained on the function-stub style, it improved at it as well.

The leftmost columns ending in "syntax" tell us how much of the Python code it generated was valid and produced no errors. You can see that the cooled version (before SFT) was already scoring pretty well at this. It could generate valid code, just not code that followed the instructions.

Even tho the syntax passing rate has not improved dramatically, the eval passes have. This indicates that it already knew how to write Python, but just didn't know how to follow instructions correctly.

What's next?

Actually, I think this is the end of Mork2.0. Although I have not gone very in depth, and could fine-tune a lot more things to squeeze more performance out of it. In the end we're quite limited by what is possible with a ~98M parameter model, trained on a laptop GPU.

All of the code for this post is in Mart100/mork-2.

As a next project, I really want to bring Mork2 to cloud computing. Renting GPUs online will allow us to increase the model size, train faster, and train on way more tokens.