Skip to content
Saran Teja Mallela
Go back

What a Data Engineer Misses About Transformers (Until You Build One)

I build data pipelines for a living. Kafka, Spark, Azure, the usual. I assumed that gave me decent intuition for how LLMs get trained. It is all just moving tensors through a DAG, right?

So over eight days in late June I worked through Karpathy’s Zero to Hero, typing every line: an autograd engine, bigram and trigram language models, an MLP with batch norm, and a GPT built up from a single attention head. I watched all ten videos and built through the GPT one. The code, commit history and all, is on GitHub.

My data engineering instincts were not missing details. They were pointed in the wrong direction, in five specific places. And then there were the bugs, which deserve their own section, because none of them crashed.

Attention is a differentiable join

Table of contents

Open Table of contents

1. Backprop is lineage tracking that carries blame

Built: micrograd, a scalar autograd engine, twice. Once following along, once from memory the same day to see what stuck.

Micrograd is about 100 lines. The core is a Value class where every value remembers what produced it: its parent values and the operation between them. If you have ever built column-level lineage for a data platform, you have already built this shape.

The difference is what flows through it. My lineage graphs answer “where did this bad row come from” once, during an incident, with a human reading the graph. Autograd answers the same question numerically, thousands of times per second, and then fixes the upstream automatically. The backward pass walks the DAG in reverse topological order and hands each node a number that means: this is how much you were responsible for the error.

def __mul__(self, other):
    out = Value(self.data * other.data, (self, other), '*')
    def _backward():
        self.grad += other.data * out.grad  
        other.grad += self.data * out.grad  
    out._backward = _backward
    return out

The += is the part worth staring at. When a node fans out to two consumers, blame arrives from both, and you sum it. Any data engineer who has debugged a fan-out DAG already has the intuition. It just runs the other way here.

Tip

If you have built data lineage, you own the mental model for autograd. Same DAG, opposite direction, and the payload is blame instead of provenance.

2. The loss function is a data quality metric

Built: bigram, trigram, and MLP language models on a dataset of 32,000 names.

Cross-entropy sounds like math-department vocabulary. From the inside it is brutally simple: the model is graded on how surprised it is by your data. The loss is the average negative log probability the model assigned to the character that actually came next.

And the number is interpretable. A loss of L means the model is effectively choosing between about e^L equally likely options at each step. The names dataset has 27 characters. Guessing uniformly gives ln(27) ≈ 3.30. Here is what each model bought me, numbers straight from my commit messages:

ModelLossEffective choices per character
Uniform guessing3.3027
Bigram (counting)2.45~11.6
Trigram (neural)2.36~10.6
MLP (dev split)2.19~8.9

Every drop in that table is structure the model found in the data. No structure, no drop. I sign off on data quality dashboards at work. Null rates, freshness, schema drift. Cross-entropy prices something none of them touch: how much learnable signal the dataset actually contains.

3. Batch size is a hyperparameter, and batch norm couples your rows

Built: MLP with activations, gradient diagnostics, and batch norm, following makemore part 3.

In my world, batch size is a throughput knob. Rows are independent. Processing 500 or 5,000 at a time changes cost and latency, never the answer. Row-level independence is something I have defended in design reviews.

Training violates both, on purpose.

First, the gradient is averaged over the batch, so batch size sets the noise level of every optimization step. Small batches give noisy gradients, large batches give smooth ones, and that noise changes what the model converges to. Batch size is not an efficiency setting. It is a hyperparameter that changes the result.

Second, batch norm. Each activation gets normalized using the mean and variance of the current batch. Your row’s forward pass depends on which other rows it happened to ship with. Rows in a batch see each other. From a data engineering purity standpoint this is horrifying. It also worked well enough to be load-bearing in deep nets for a decade.

The train/eval distinction exists because of this coupling. At inference there is no batch to normalize against, so the model uses running statistics collected during training. Forget model.eval() and your predictions quietly depend on which rows shipped together. Quietly is the theme of this post. Hold that thought for section 5.

4. Attention is a differentiable JOIN

Built: a GPT from a single attention head up to a 6-layer block stack, trained at reduced scale on my laptop.

This was the part I expected to be alien. Instead I recognized it.

Every token emits a query vector. Every token emits a key and a value. You dot every query against every key. That is an all-pairs comparison: a cross join. The causal mask is a join condition on position. Softmax turns raw scores into weights that sum to one. Then you take the weighted sum of the values, which is an aggregation.

Here is attention written as the query it secretly is:

SELECT q.position,
       SUM( softmax(dot(q.vec, k.vec) / sqrt(d_k)) * v.vec )
FROM   tokens q
JOIN   tokens k  ON k.position <= q.position   -- causal mask
JOIN   tokens v  ON v.position =  k.position
GROUP  BY q.position

And here is the real thing, from my repo:

k = self.key(x)
q = self.query(x)
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5
wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
wei = F.softmax(wei, dim=-1)
out = wei @ self.value(x)

It is a join where the ON clause is learned, matches are soft instead of boolean, and the aggregation is a weighted average. The 1/√d_k scaling keeps dot products from saturating softmax. I think of it as normalizing the join keys so no key wins on magnitude alone. And the KV cache everyone discusses at inference time is a materialized view of the join’s right-hand side.

The analogy is load-bearing but not perfect, so let me break it before someone in the comments does: SQL joins do not backpropagate through their predicate. Attention does. The join condition itself is what gets trained. That difference is the whole point.

5. Deep learning bugs don’t crash. They converge.

That clean attention code above? It did not start that clean. While drafting this post I re-read my Head.forward and found three bugs. All three trained fine. Loss went down. Generated text got less random. The model was wrong the entire time.

k = self.query(x)
k = self.key(x)
q = self.query(x)
wei = q @ k.transpose(-2, -1) * C**-0.5
wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5

Bug one: k was computed from the query projection. My key matrix was defined, initialized, and never used. Queries were attending to other queries. The model still learned, because query-times-query is still a valid similarity. Just not the one the architecture intends.

Bug two: the scaling constant used the embedding dimension (384) instead of the head dimension (64). My attention logits were about 2.4x smaller than designed. Softmax ran flatter than intended. Nothing complained.

Bug three: my block stack hardcoded num_heads=4 while the config said n_heads = 6. The config was silently ignored. Every data engineer who has debugged config drift just felt something.

Here is the part that rewired me. During the same build, I hit a fourth bug: a shape mismatch when I introduced the embedding dimension, where the embedding table stopped producing logits and cross-entropy got tensors of the wrong width. That one crashed instantly, and it was fixed within the hour (commit e1ade3c, the fix is adding the lm_head projection). The bug that crashed was the cheapest one. The three that mattered never crashed at all.

In data engineering, wrong code throws. In deep learning, wrong code converges. A pipeline with a bad join key produces obviously broken output. A transformer with a bad join key produces a slightly worse language model, and no dashboard on earth flags it. The only defenses are the ones Karpathy models throughout the series: read your shapes out loud, verify your initial loss against theory (ln of vocab size, and yes, mine matched at 4.17 for a 65-character vocab), and re-read your forward pass like it owes you money.

Warning

If your model trains without errors, you have verified nothing except that the shapes broadcast. Broadcasting is not correctness. Sometimes it is the opposite.

The one I haven’t built yet

Full disclosure: the last video, reproducing GPT-2 (124M), is the one artifact I studied but have not built. Four hours of systems engineering, and I watched it like a systems engineer: pausing, taking notes on the GPU tradeoffs, and running the numbers myself. The build is next, and it gets its own post with real throughput numbers.

But the napkin math alone rearranged my head. In data engineering, data size is the size of the data. In training, the working set is a multiple of the artifact before you process a single row. GPT-2 124M in fp32:

Where a 124M-parameter model actually lives

The optimizer alone is twice the model. Adam keeps two running statistics per parameter, a first and second moment, which a data engineer immediately recognizes as 124 million tiny streaming aggregations that never stop. Gradient accumulation, the standard trick when a batch does not fit, is a windowed aggregation: run micro-batches, sum gradients, step once. I have written that exact pattern in Spark.

My favorite detail from the video is pure data engineering: GPT-2’s vocabulary is 50,257 tokens, an ugly number for GPU kernels, so you pad it to 50,304, which is divisible by 128. Nicer memory alignment, faster kernels, a few thousand embedding rows that exist purely so the hardware stays happy. We pad partitions for the same reason. Nobody tells you the frontier of AI includes making numbers divisible by 128.

What building it changed

I work on evaluation infrastructure for clinical LLMs, and this rebuilt how I read my own tooling. Loss curves read as data quality reports now. Attention maps read as query plans. And section 5 is why I believe evals are not optional: when wrong code converges, measurement is the only thing standing between you and confident nonsense.

If you move data for a living, you already own most of the mental models. Lineage, joins, aggregations, streaming state, batch semantics. They are just pointed the wrong way. Spend a week and turn them around.

The code is here, bugs, fixes, commit history and all. Next post: the GPT-2 reproduction, with throughput numbers. After that: turning a clinical eval into an RL environment with verifiable rewards.


Share this post:

Next Post
Why I'm Spending 8 Weeks Building Toward Verifiable Clinical AI Environments