Train Tiny GPT End-to-End

Implement train_tiny_gpt: train a small causal language model over a supplied sequence of batches. Each step predicts next tokens, calculates gradients by hand and updates the model with SGD. Carry those updated weights into the next step.

The order matters. This is a sequence of updates, not one gradient computed over the whole corpus and not a collection of independent updates from the initial weights. There is no shuffling, random sampling or generation loop.

Inputs and output

• tokens_corpus: (n_steps, N, T+1), integer-valued IDs in [0, V), delivered as floats. Batch k is tokens_corpus[k].
• w_emb: (V, d), token embeddings.
• pos_embed: (T, d), position embeddings shared across examples.
• blocks_weights: (L, 6, d, d), each block holding [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2].
• w_head: (d, V), a separate vocabulary projection.
• num_heads: a positive integer dividing d.
• lr: a finite, nonnegative learning rate, constant across steps.
• n_steps: an integer from 0 through 5, matching the corpus's first dimension.

N, T, V and d are positive; L may be zero. Parameters are finite and share a floating-point dtype and device; the corpus is on that device. Shapes agree. Leave all inputs unchanged and preserve parameter dtype and device.

After all steps, flatten the final weights in row-major order and concatenate [w_emb, pos_embed, blocks_weights, w_head]. Return this one-dimensional tensor of length Vd + Td + L6dd + dV. Do not return intermediate weights or losses. With n_steps=0, the corpus has shape (0, N, T+1) and no batch is read. With lr=0, all updates are zero. In both cases, return the original parameters packed in the specified order.

Sequential updates

current = supplied parameter tensors
for k = 0, ..., n_steps-1:
    tokens = tokens_corpus[k]
    gradients = next_token_loss_gradients(tokens, current)
    current = current - lr * gradients
return flatten_and_concatenate(current)

Recompute the forward pass and gradients using the current weights at every step. Within a step, compute all gradients before updating any parameter. If you reuse a helper that returns packed weights, unpack them into their original shapes before the next step; a flat tensor is not a tuple of four tensors. The helper from another exercise is not supplied automatically.

For a batch row [2,4,1,3], model inputs are [2,4,1] and targets are [4,1,3]. There are N*T next-token predictions in each step. The following sections define that step completely.

1. Forward pass

Set input_ids = tokens[:, :-1] and targets = tokens[:, 1:]. Start with x = w_emb[input_ids] + pos_embed, broadcasting positions across the batch. Apply the L blocks in order using post-normalization:

q, k, v = x @ w_q, x @ w_k, x @ w_v
y = LN(x + causal_attention(q, k, v) @ w_o)
x = LN(y + GELU(y @ w_mlp1) @ w_mlp2)

For attention, split projected features into consecutive heads of width d/num_heads. Divide q @ k.T by sqrt(d/num_heads), mask future keys with negative infinity, then apply stable softmax over key positions. Query t may read keys 0 through t, including itself. Concatenate heads in their original order. Different batch examples never attend to each other.

LN uses each token's feature mean and population variance: (token - mean) / sqrt(variance + 1e-5). There is no learned scale or bias. Use GELU(u) = 0.5u(1 + tanh(sqrt(2/pi)(u + 0.044715u**3))). The feed-forward width is d. There are no biases, dropout, extra embedding scaling or additional final normalization. No token ID is treated as padding.

Compute logits = x @ w_head. With L=0, project token-plus-position embeddings directly. Keep the intermediate values needed to differentiate each block.

2. Loss and backward pass

Use cross-entropy against targets at all NT model positions, averaging once over NT. For one vocabulary row, its logit gradient is (softmax(logits) - one_hot(target)) / (N*T). There is no additional averaging when accumulating parameter gradients. Use stable log-softmax if calculating the loss value; the return does not include it.

Differentiate the output head, then visit the blocks in reverse order. For post-normalization, first pass the gradient through the layer normalization, then send it through both the residual skip and the sublayer. Each block's normalization derivative must use that block's saved normalized output and standard deviation. The final block's output cannot stand in for earlier ones.

Flatten the batch and position axes when forming gradients for shared matrices. For example, dw_head = x_flat.T @ dlogits_flat. Future attention positions have zero probability and zero score gradient.

Sum position-embedding gradients across the batch. Accumulate token-embedding gradients for every occurrence of each input ID; repeated IDs add to the same row. Only tokens[:, :-1] are embedding lookups. A token occurring solely in the final target position does not directly update its input embedding row, although its target class affects gradients elsewhere through cross-entropy.

3. Simultaneous SGD update

Compute every gradient from the weights at the start of this step, then apply new_parameter = parameter - lr * gradient to all four parameter groups. There is no clipping, weight decay, momentum or optimizer state.

Implement the forward operations and derivatives from basic tensor operations in either framework. Do not use autograd, backward(), jax.grad, packaged transformer layers or optimizer helpers. The same manual-gradient requirement applies to both variants.