Train Causal LM Pretraining Step

Implement train_causal_lm_step: predict the next token at each input position and take one SGD step using gradients you calculate by hand.

Each supplied sequence contains one more token than the model processes. For tokens [2, 4, 1, 3], feed [2, 4, 1] into the model and use [4, 1, 3] as targets. The final token is a target only. Causal attention prevents an input position from reading later inputs, including the next token it must predict.

Inputs and return value

• tokens: (N, T+1), integer-valued IDs in [0, V), delivered as floats.
• w_emb: (V, d), token embeddings.
• pos_embed: (T, d), position embeddings shared across examples.
• blocks_weights: (L, 6, d, d), transformer blocks, each holding [w_q, w_k, w_v, w_o, w_mlp1, w_mlp2].
• w_head: (d, V), a separate output projection, not tied to w_emb.
• num_heads: a positive integer dividing d.
• lr: a finite, nonnegative learning rate.

N, T, V and d are positive; L may be zero. Parameters are finite and share a floating-point dtype and device, with tokens on that device. Shapes agree. Cast token IDs to integers and leave every input unchanged.

Return a one-dimensional tensor of updated weights, flattening each in row-major order and concatenating [w_emb, pos_embed, blocks_weights, w_head]. Its length is Vd + Td + L6dd + dV. Preserve parameter dtype and device. Return weights, not gradients or a loss. With lr=0, return the original weights in that same order.

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 original weights, 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.