GRU-L3 V3a: Gradient Checkpointing, 75M Parameters, and the Path Past TF 0.75

The Memory Wall

Our previous GRU-L3 V3 model was a 4.8M parameter character-level RNN trained at D=512, B=64, T=1024 on a single RTX 3080 8GB. It reached 74% top-1 eval accuracy in 6 epochs with teacher forcing decayed from 1.0 to 0.75. But it plateaued hard — eval loss degraded every epoch after TF dropped below 0.95. The model didn't have enough capacity to model the free-running generation distribution.

The obvious fix was more parameters. But scaling from D=512 to D=2048 meant 16x more parameters (75.9M) and a VRAM footprint that blew past 8GB at the batch sizes we needed. The activation buffers alone — h_normed, gate_cache, combined_input — scaled linearly with batch size, sequence length, hidden dim, and layer count. Something had to give.

A18: Segment-Wise Gradient Checkpointing

The solution was gradient checkpointing — trading compute for memory by recomputing forward activations during the backward pass instead of storing them all. We call it A18 in our optimization plan, and it's the hardest low-level GPU work we've done on the cat_chat stack.

The core idea: divide the sequence dimension T into K segments of size S = ceil(T/K). Instead of allocating activation buffers for the full sequence, allocate them for one segment. During backward, recompute the forward pass for each segment before consuming its gradients.

The implementation required careful offset management. Each timestep within a segment uses t_local = t % ckpt_S to index into the smaller buffers. The recompute lambda replicates the full forward GRU pipeline — reading h_normed and combined_input from full-T buffers, writing checkpointed buffers at t_local offsets. The backward pass processes segments in reverse order, recomputing each segment's forward activations before consuming them.

One buffer remains stubbornly full-T: h_normed. The classifier backward and cross-segment hidden carry need all timesteps. Segmenting it would require a backward redesign for the loss path — a future optimization.

The VRAM Math

At B=64, K=4, the savings were dramatic:

The 44% VRAM reduction (7749 → 4322 MiB at B=64) came from shrinking five checkpointed buffers from T to S: gate_cache, h_raw, rmsnorm_in, rmsnorm_sq, and r_h_prev. At K=4, S=128 — a quarter of the original T=512.

With that headroom, we doubled the batch size to B=128. The VRAM climbed to 7064 MiB — still under the 8192 MiB ceiling with 1.1 GB of safety margin. The batch time increased from 9.5s to 15.6s, but we process 2x more sequences per batch. Net throughput: ~4200 tokens/second vs ~3400 tokens/second at B=64 — a 22% throughput improvement with 2x the gradient stability.

Scaling Decisions: Width, Depth, Batch, Context

With freed VRAM, the question was: spend it on width (D), depth (L), batch (B), or sequence length (T)? We researched each option.

Gradient Clip Scaling

The V3 model used clip_norm = 1.0 — the de-facto standard for Adam-type optimizers, backed by arXiv 2406.04443. But V3 had 4.8M parameters. V3a has 75.9M — 16x more. The global L2 gradient norm scales with sqrt(num_params), so the same clip threshold is effectively 4x more aggressive on the larger model.

At clip_norm = 1.0, V3a's early gradients (gn_pre ~4.5) were being scaled by 0.22 — cutting 78% of the gradient signal. At clip_norm = 2.5, the same gradients scale by 0.56 — still controlled, but with 2.5x more update magnitude. The sqrt-scaled value would be 4.0, so 2.5 is conservative.

The 20-Epoch TF Decay Schedule

V3 Run 2 peaked at epoch 2 (TF=0.95, eval top1=74.04%) then degraded steadily as TF dropped. By epoch 6 (TF=0.75), eval top1 had fallen to 73.00%. The model couldn't maintain quality when weaned off teacher forcing.

V3a's 16x larger capacity may delay that plateau. More parameters means more capacity to model the free-running distribution, not just the teacher-forced one. The schedule:

Each epoch runs as a separate 1-epoch session with resume = true, loading the previous checkpoint. This lets us monitor eval metrics between epochs and halt if degradation appears earlier than expected. At 7.7 hours per epoch, the full 20-epoch run takes ~6.4 days.

Epoch 0 Telemetry: First Light

The E0 run is in progress as of this post. Early telemetry (first 50 steps, 2.8% of epoch):

Loss is descending from the expected log(100) = 4.605 starting point. BPB dropped from 6.64 to 3.03 in 49 steps. The gradient norm settled below the clip threshold by step 41 — a sign that clip_norm = 2.5 is well-calibrated for this model size. No NaN, no Inf, no anomalies.

For comparison, V3 at the same step count was still at BPB ~4.6. V3a descends faster per optimizer step because B=128 sees 2x more tokens per batch, giving each gradient update more data to work with. The larger model's higher capacity is already showing in the steeper BPB descent.

The Data Question

The telemetry from this run serves a dual purpose: training the model and measuring the data gap. If the model peaks early and degrades, that's a data problem, not an architecture problem. The fix is more dialogue data, not more parameters. If it holds steady through the TF decay, the architecture is working and we can focus on data collection for the next scale-up.

What's Next

The immediate roadmap after E0 completes:

1. E1-E20 incremental resume with TF decay -0.05/epoch, monitoring eval at each boundary
2. Parity control on the release build to validate A18 checkpointing end-to-end
3. Data collection — if overfitting appears before E5, we need 500M-1.5B tokens to reach Chinchilla optimal
4. h_normed segmentation — the last full-T buffer, requiring a classifier backward redesign
5. FP16 training — halve weights/grads/optimizer VRAM, enabling D=3072 or B=256

The A18 checkpointing implementation, VRAM budgeting, and clip scaling are all implemented and running. The 20-epoch TF decay schedule and data gap analysis are in progress. The h_normed segmentation and FP16 training are roadmap.