How Fixing an OOM Error Quietly Broke My Model

For the past few weeks I have been hand-writing a simple transformer from scratch — building the pieces myself straight out of the original Attention Is All You Need paper rather than reaching for a library — and training it on the Parenting Stack Exchange dump. The goal was to learn by doing; what I actually learned was how many quiet ways a from-scratch model can fail.

After parsing Posts.xml, the corpus came out to:

• 6,872 questions (6,734 with at least one answer)
• 22,601 answers → 22,601 question/answer pairs (avg 3.29 answers per question)
• ~60.3M characters once flattened into a single training stream
• ~15.1M word-level tokens after tokenization

Each pair is serialized as <Q>…<A>… and joined with a <D> document separator, then fed to the model as one long autoregressive stream. The <pad> token is used for batching, and anything outside the 300 most-common words becomes <unk>.

The architecture

It is a deliberately small, GPT-style decoder with weight tying — the output projection reuses the embedding matrix (Eᵀ) instead of learning its own:

The full configuration:

Group	Hyperparameter	Value
Data	MAX_VOCAB	300 (+2 reserved)
Data	sequence_length	254
Model	embedding_dimensions	200
Model	num_blocks	1
Model	feedforward_hidden_layer_depth¹	1 (used as width)
Training	batch_size	1000
Training	learning_rate	0.01
Training	EPOCHS	6

Training ran on Kaggle’s 2× T4 GPUs — 16 GB of memory each.

¹ This hyperparameter was meant to be width, not depth. It forced the feed-forward hidden layer to 1 neuron. That bug is the second half of the story.

Bug 1: the fix that caused the real bug

After fixing bugs, my model finally started to train, but before any diagnostic from the first batch came in, Kaggle threw an out-of-memory error on me. I eventually singled it out to unbounded vocabulary size.

Essentially, I was using word-level tokenization, so every distinct word in the corpus became its own token — the vocabulary grew into the tens of thousands. The output layer materializes one logit per vocab entry per position, so the logits tensor alone is:

logit bytes = batch_size × sequence_length × vocab_size × 4 bytes (fp32)

At V = 300:
= 1000 × 254 × 300 × 4 bytes
= 304,800,000 bytes
≈ 305 MB

That is ~305 MB just for the logits — and the softmax probabilities plus backward-pass gradients roughly triple it. Scale V into the tens of thousands and the byte count grows linearly into the tens of gigabytes, sailing straight past the 16 GB ceiling.

Seeing this, I capped the vocabulary size to just the 300 most common tokens and mapped every other word to <unk>.

Why V=300 ended up being a disaster

The model did train. When I checked back in the morning the test accuracy and loss were 13.39% and 3.7093 respectively. Excited to see the toddler give me parenting advice, I threw in the prompt “How to raise” at temperature 0.1 and got back <unk>. Raising the temperature to 0.7 mixed in tokens like the, but the output was still predominantly <unk>.

I realized the bound of 300 might be too harsh — most tokens were now <unk>. So I averaged the model’s output probabilities over 10 test batches and laid them side by side with the corpus token frequencies. The model had not learned any Q&A structure; it had simply memorized the unigram distribution of the corpus. They don’t just correlate — all 20 of the top-20 tokens appear in the exact same rank order (a 20/20 match), and the magnitudes line up to within a few hundredths of a percent:

Bug 2: a feed-forward layer one neuron wide

The vocabulary cap was the most visible failure, but it wasn’t the only one. When I went back and audited the notebook, I discovered I had accidentally made the feed-forward hidden layer only one neuron wide.

Every 200-dimensional token embedding was funneled through this single scalar and back out again.

I had a hyperparameter called feedforward_hidden_layer_depth and set it to 1, but I was using it as the width of the hidden layer:

self.feed_forward = nn.Sequential(
    nn.Linear(embedding_dimensions, feedforward_hidden_layer_depth),  # 200 -> 1
    nn.ReLU(),
    nn.Linear(feedforward_hidden_layer_depth, embedding_dimensions),  # 1 -> 200
)

A proper transformer feed-forward network is usually the widest part of the block — about 4 × d_model (the original paper uses d_model = 512 and d_ff = 2048), or ~800 neurons here. It holds most of the model’s parameters and most of its representational capacity.

The parameter count comes straight from the two nn.Linear layers:

Linear(200, 800) weights: 200 × 800 = 160,000
Linear(200, 800) biases:            800 =     800
Linear(800, 200) weights: 800 × 200 = 160,000
Linear(800, 200) biases:            200 =     200
-----------------------------------------------
Proper FFN:                      321,000

My FFN:
Linear(200, 1) weights:   200 × 1   =     200
Linear(200, 1) biases:              1 =       1
Linear(1, 200) weights:   1 × 200   =     200
Linear(1, 200) biases:            200 =     200
-----------------------------------------------
Actual FFN:                          601

So the correct FFN should have been roughly 321,000 parameters (200 → 800 → 200 plus biases); mine had 601 — a ~530× cut to the single most important sublayer. No vocabulary fix was probably going to rescue a block that can’t hold more than a single number of internal state.

Takeaways

a) The metrics being optimized are a proxy for the model quality. Always play with your model’s outputs. For example: 0-1 accuracy and loss of 13.39% and 3.7093 respectively on 300+ vocabulary seemed promising yet the model is useless.

b) I need to learn about the tradeoffs of different tokenization schemes. I will go one step further and say I should invest time learning about hyperparameter tradeoffs too.