Pet project write-up

GPT-2, from scratch, learns to call functions

by Mansur Orazkhan · PyTorch, no transformers, no PEFT · source on GitHub · try it live

I wrote GPT-2 in raw PyTorch — attention, layer norm, the training loop, everything — and then fine-tuned it until it stopped writing internet text and started emitting machine-readable function calls. This page is the story of my adventure into the world of Large Language Models: my tenth (or maybe more) introduction to Transformers, and the training and fine-tuning of a GPT model from scratch.

Real output from the fine-tuned 355M model, greedy decoding, on a schema + request it never saw in training.

Where this came from

I came across Sebastian Raschka's book Build a Large Language Model From Scratch, which teaches everything about LLMs. I even made a LinkedIn post where I said that I would finish this book and return with a review or some cool information.

Overall, this is a great book that takes you from the most basic stuff to the very final steps of LLM training and fine-tuning. It has a great narrative, clear structure, and practical code examples that show you everything in detail, step by step.

After practicing with the book, trying to recreate each code block of the transformer and GPT training, I wanted my own twist. Due to my background in multi-agent systems, I wanted to teach an LLM to call functions. Function calling is what makes LLMs useful as agents — the model reads a JSON schema, decides whether a tool is needed, and produces a structured call instead of prose.

I did not want to spend a lot of resources training big models, so I settled on GPT-2 (355M), which is small enough to fine-tune on a single GPU in a few hours. So the question was:

Can a 2019 GPT-2, which I loaded from OpenAI's checkpoints, learn that protocol?

That extension — dataset, training run, evaluation harness — was my initiative, and it's the fun part.

Part 1 — The transformer, hand-built

No AutoModel. The whole model is ~200 lines of PyTorch: token + positional embeddings, a stack of pre-norm transformer blocks, a final layer norm, and a linear head that scores all 50,257 tokens of the vocabulary. Watch a forward pass:

Token + positional embeddings 50257×1024 · 1024×1024
× 24 blocks (355M)
LayerNorm → Masked multi-head attention → +residual 16 heads
LayerNorm → FeedForward (GELU) → +residual 1024→4096→1024
Final LayerNorm → LM head → logits over the vocabulary 1024×50257

One forward pass: six tokens go in, a probability for every possible next token comes out. Generation is just doing this in a loop.

The causal mask

The one rule that makes a language model a language model: a token may attend to itself and the past, never the future. In the attention matrix that's a triangle — everything above the diagonal is set to −∞ before the softmax:

Each row is a token computing attention; darker blue = more weight. The dashed cells are the future — masked out, always.

Multi-head attention, the GELU feed-forward, layer norm — each is a small, testable module (src/gpt2fc/model/), covered by unit tests that check the causal mask actually blocks information flow and the parameter count matches the analytic formula (354,749,440 for 355M — you learn to love exact numbers when a single transposed weight silently breaks everything).

Part 2 — Teaching it to call functions

The dataset is Glaive Function Calling v2: 112,960 dialogs, each with a function schema and a conversation where the assistant sometimes calls the function. A training sample looks like this:

###SYSTEM: You are a helpful assistant with access to the following functions…
{ "name": "get_current_weather", "parameters": { "location": … } }
###USER: What's the weather like in Almaty right now?
###ASSISTANT: <functioncall> {"name": "get_current_weather",
  "arguments": '{"location": "Almaty"}'}

The trick: only grade the answer

If you compute the loss over the whole sequence, the model wastes capacity learning to predict the prompt it was just given. Instead, the collate function shifts the targets by one and stamps ignore_index=-100 on every prompt and padding token — cross-entropy only flows through the assistant's response:

masked, −100 — no gradient loss computed — this is what the model learns

The run

Model
355M
GPT-2 medium, full fine-tune
Training samples
112,960
1 epoch, batch size 8
Hardware
1× T4
Kaggle free tier
Wall clock
9.45 h
AdamW, lr 5e-5
Cross-entropy loss on response tokens train validation
Validation loss falls from 1.85 to 0.135 over 12,625 steps. Hover the chart for exact values.
View as table
StepTrain lossVal loss

That cliff at the start is the model learning the format — the ###ASSISTANT: convention, the <functioncall> tag, the JSON shape — within a few hundred steps. The long flat tail is the harder part: learning when to call, and how to fill in the arguments.

Before / after

Same prompt, same greedy decoding. Only the weights differ.

Base GPT-2 355M — pretrained only

Fine-tuned 355M

The base model is a web-text predictor: it treats the prompt as a document to continue. Here it invents a schema-shaped JSON blob of its own, then hallucinates a ###SYSTEM: turn — pure surface imitation of the prompt's format. On other prompts it parrots the schema back or just pads whitespace until the token budget runs out. The failure mode varies, but the constant is that it never makes a call: it has no concept of the dialog, the tool, or stopping. One epoch of fine-tuning installs the whole protocol: answer as the assistant, treat the schema as callable, pull "almaty" out of the user's sentence, emit end-of-text, done.

Don't take my word for it — both models run live in a Hugging Face Space: type your own request, edit the schema, and watch them stream side by side.

The bug that made a working model look broken

My first evaluation said the model failed at function calling almost completely — single-digit accuracy. The generations looked right, but nothing parsed. The culprit wasn't the model. Glaive's ground truth (and therefore my faithfully-imitating model) emits almost-JSON:

{"name": "get_current_weather", "arguments": '{"location": "Almaty"}'}

The arguments object is wrapped in single quotes — invalid JSON, so json.loads throws. The fix: find the payload span by counting brace depth (regexes die on nested objects), then unwrap the single-quoted blob before parsing. Measured on the ground-truth answers themselves:

naive json.loads0%
brace-depth + single-quote unwrap0%

Share of the test split's 1,856 ground-truth function calls each parser can read.

Lesson: evaluation infrastructure deserves the same paranoia as model code. I nearly wrote off ~10 hours of GPU time over a quotation mark.

Results

Evaluated on the full held-out test split: 3,929 dialogs the model never saw. In 1,856 of them the correct answer is a function call — those are the denominator for all four metrics below. Scoring is strict: if a prediction doesn't parse as JSON, it counts as wrong in every metric.

Parseable call produced
0%
1,641 of 1,856 expected calls
Correct function name
0%
1,633 — nearly every parsed call
Correct argument keys
0%
1,510 with the exact key set
Exact match, all values
0%
1,439 perfect calls, end to end