How much can a
small model learn?

01 / THE QUESTION

A useful assistant within a real memory budget.

The long-term goal is a compact assistant that can explain ideas, follow instructions, work with supplied information, and eventually call a small set of tools. A search tool could bring in evidence; a Python tool could handle arithmetic or string operations. Those are future capabilities to train and evaluate. The current run is building the underlying language model.

Pretraining teaches next-token prediction across ordinary text and code. It does not automatically teach a model when to ask for clarification, how to use search, or how to stop repeating itself. Those behaviors need carefully chosen examples, post-training, and tests of actual responses.

The central research question is practical: with a limited parameter and compute budget, which choices produce the most useful behavior? Smaller models make the tradeoffs especially visible. Extra text is only helpful if the model can learn something useful from it.

02 / INSIDE THE MODEL

One token in. A prediction for the next.

Text becomes token IDs. An embedding turns each ID into a vector. Thirty-two transformer blocks update those vectors using the preceding context. The output projection produces scores over the vocabulary, and training rewards the score assigned to the actual next token.

1. Text → tokenssteinshark_32k BPE
2. Embedding1,024 dimensions
3. 32 transformer blocksCausal attention + feed-forward
4. Vocabulary scoresTied embedding weights

Fitting the work into memory

The model uses activation checkpointing: some intermediate results are recomputed during backpropagation instead of being kept in memory. The vocabulary projection and loss are also evaluated in small chunks. Both choices reduce memory pressure, but recomputation costs time. Memory saved is not free throughput.

The current optimizer is AdamW with 8-bit optimizer states. Computation uses BF16 while the stored model parameters remain FP32. These are separate choices: an 8-bit optimizer does not mean all model weights and activations occupy eight bits.

03 / WHAT IT READS

The dataset is part of the model.

The current sampling plan assigns 75% to FineWeb-Edu, 20% to cleaned and relevance-filtered Wikipedia, and 5% to filtered Python from NPset. The two FineWeb collections are sampled separately at 48% and 27%. These are sampling weights, not guarantees that the model will encounter each document exactly once.

Wikipedia: remove the scaffolding, keep the meaning

References, navigation templates, list-only pages, and duplicate text can consume a small model’s context without providing much continuous explanation. The cleaner removes much of that scaffolding, while a later relevance pass reduces selected narrow article categories and protects a core article list.

Cleaning needs inspection. Removing every template also removed pronunciation symbols from the article about the letter A. Removing a sentence because it contained markup once damaged the opening definition of Aircraft. Those failures led to targeted fixes and regression examples. “Cleaner” text is not necessarily more faithful text.

Python: learn manageable patterns

The Python selection emphasizes small, understandable programs: basic imports, file and text handling, and simple data structures. Source files and extracted functions need different filters. An isolated function can depend on an import or global variable that is no longer present, even when the function parses correctly.

The tokenized library is stored in Parquet shards with source IDs, document IDs, token IDs, and checksums. It preserves document boundaries; the training loader inserts end-of-document tokens when assembling sequences. Keeping those boundaries matters for reproducibility and for understanding exactly what the model saw.

04 / TRY THE MECHANICS

A few lines make the idea concrete.

These are teaching examples, simplified from the training pipeline. They run on small inputs; they do not train or query the model in your browser.

def training_windows(documents, context, eos_id):
    buffer = []
    for document in documents:
        buffer.extend(document)
        buffer.append(eos_id)
        while len(buffer) > context:
            inputs = buffer[:context]
            targets = buffer[1:context + 1]
            yield inputs, targets
            # Keep the next token as the next window's first input.
            del buffer[:context]

documents = [[11, 12, 13], [21, 22, 23]]
for inputs, targets in training_windows(documents, 4, 99):
    print("input:", inputs, "target:", targets)
# input: [11, 12, 13, 99] target: [12, 13, 99, 21]
# The final short remainder is omitted in this toy example.

2. Trade recomputation for memory

This is the core pattern used around a transformer block. PyTorch can recompute intermediate activations during the backward pass. Saving a checkpoint to disk is a different operation, despite the shared name.

from torch.utils.checkpoint import checkpoint

# x: hidden states; cos/sin: rotary position information
x = checkpoint(block, x, cos, sin, use_reentrant=False)

# Later, loss.backward() recomputes needed intermediates.
# Gradients still train the block's parameters.

3. Make the time budget visible

At billions of tokens, small changes in sustained throughput become days of training. This calculator assumes a constant rate and excludes validation, checkpoint writes, downtime, and context-phase changes.

05 / WHAT WENT WRONG

The failures are useful measurements.

More layers did not settle the question

An earlier model expanded from 22 to 32 layers and continued pretraining. Validation loss improved modestly. Extra capacity was not evidence by itself of better instruction following. Comparing models needs matched prompts, evaluation settings, and enough examples to separate a real change from sampling noise.

A fluent sentence can still be a poor answer

Early supervised fine-tuning produced recognizable assistant responses, but also repeated list items, vague advice, and weaker behavior after a few conversation turns. A no-repeat n-gram setting can block some exact repetition at generation time. It does not prove the model has learned to make distinct, useful points.

Low throughput needs a controlled comparison

A memory-saving model can train much more slowly than an older implementation with similar parameter count. Context, microbatch size, compilation, loss computation, and checkpointing all affect the comparison. A useful benchmark holds those settings fixed and measures enough steps to get beyond warmup.

06 / WHAT COMES NEXT

Pretrain. Measure. Then teach the interaction.

The immediate work is to complete the new pretraining run and evaluate the resulting checkpoint. The next phase will use reviewed conversations to teach direct answers, clear formatting, grounded responses, and continuity across turns.

Search and restricted Python execution are research directions. A tool-capable assistant must learn when a tool is appropriate, how to construct a valid request, and how to use the returned result. It also needs a separate execution environment with appropriate restrictions. Tool syntax alone provides none of those guarantees.

07 / KEEP READING

Follow the ideas back to their sources.