Tot-360M's First Real Training Run: 405 Million Tokens on a Kaggle TPU

Tot-360M is the small language model I'm training from scratch to sit inside Tot. Two weeks ago the honest status was that not one training step had worked on real hardware. This week a run finished. Here is what happened, with the numbers from the logs and nothing rounded up.

the short version

• Model: 345,359,296 parameters. 32 layers, a 960-wide hidden size, grouped-query attention with 15 query heads and 5 key/value heads, and a 32,000-token tokenizer built for this project.
• Hardware: a free Kaggle TPU v5e-8, which is eight TPU cores working on one model.
• Run length: 8 hours 37 minutes of training.
• Steps: from step 439 to step 24,785.
• Tokens: 398,495,328 new tokens this run, 405,680,880 in total.
• Speed: about 18,100 tokens per second while actually training.
• Loss: 5.19 at the start of the run, 2.67 on the last logged step.

That last number needs a caveat, so it gets its own section below.

three failures first

The run that worked was the fourth attempt on this notebook. The first two died during setup, at 395 and 778 seconds. The third died at 288 seconds, and that one taught me something about Kaggle.

Kaggle mounts the outputs of other notebooks through a cloud file system. On a TPU machine, some of those folders don't really exist until you ask for them by exact path. My code searched for the training data with a recursive search, and the search walked straight past folders that hadn't been touched yet. It found nothing and gave up.

The fix was two things. First, probe each known path one folder at a time so the mount wakes up. Second, copy the data and the 4 GB starting checkpoint onto the machine's local disk once. Before that, all eight TPU workers were each pulling the same 4 GB file over the network at the same moment.

what the loss number means

Training loss is a measure of surprise. At every position in a document the model gives a probability to every token in its vocabulary, and the loss is based on how much probability it gave the token that actually came next. Lower is better.

A loss of 5.19 means the right answer got, on average, about a 0.6% share. A loss of 2.67 means about 7%. Put another way, the model went from being torn between roughly 180 equally likely tokens to roughly 14.

The caveat is that 2.67 is a single step, not an average. Individual steps near the end swung between about 0.5 and 4.1, depending on whether the batch was predictable code or messy prose. The trend is clearly down, but this is training loss on the training data. It is not an evaluation and not a benchmark, and I'm not going to quote it as one.

why saving was so slow

The trainer saved a checkpoint every 250 steps, which came to 99 saves over the run. Each one is 4.07 GB, even though the model itself is only about 1.38 GB.

The extra space is the optimizer. The training algorithm (AdamW) keeps two running averages for every single parameter. One tracks which way each weight has been moving, and the other tracks how jumpy its updates have been. Both are needed to resume without a bump, so a full checkpoint is three model-sized tables: the weights plus those two averages.

Writing 4 GB to a network disk took around 90 seconds each time. Ninety-nine of those cost about two and a half hours, which is why the end-to-end rate was about 12,800 tokens per second even though the training itself ran at about 18,100.

mixed precision, briefly

The heavy matrix maths runs in bfloat16, a 16-bit number format the TPU is built around. The master copy of the weights stays in 32-bit floats. A single update nudges a weight by something like a ten-thousandth, and bfloat16 is too coarse to hold a change that small. The weight would round straight back to where it started and the model would stop learning.

what changes for the next run

The first data pack was about 43% code, and the model was already predicting code well, with some code batches near a loss of 0.5. So I rebuilt the second pack around language. It has 626,642,944 tokens:

• 321 million tokens of FineWeb-Edu, educational web text filtered for quality
• 193 million tokens of English Wikipedia, with Formula 1 articles repeated so the model sees more of them
• 97 million tokens of MIT and Apache licensed code
• 13 million tokens of Korean
• a small curated set covering everyday English slang, Formula 1 strategy and history, casual Korean, and OB Studio's own projects

The next run also saves every 1,000 steps instead of every 250. That should give back close to two hours of training in the same session. My estimate is about 518 million more tokens from one session, which would leave a short third run to finish the pack. That is an estimate, not a result.

what a .gguf file is

When the next run finishes, it will also write the model out as a .gguf file. GGUF is a single-file format used by llama.cpp, Ollama and LM Studio. It holds the weights plus everything needed to run them: the architecture settings, the tokenizer vocabulary and the special tokens.

It leaves out the optimizer averages, so the 16-bit version is about 690 MB instead of 4.07 GB. A 4-bit version for laptops should be roughly 240 MB. I tested the whole chain on a miniature version of the model: resume, train, export, then load and generate in llama.cpp. The full-size file hasn't been run yet.

where it actually stands

A trained checkpoint exists now, which wasn't true two weeks ago. It is partway through pretraining. It hasn't been evaluated, and it isn't published. It also hasn't had any personality or conversation training yet, so it can continue text but isn't a chat model.

The current status lives on the Tot-360M page, and the app it is being built for is Tot.