Backward Bug Hunts and Vulkan Tiled GEMM: Pi-Train LoRA SFT Hardening

The symptom: NaN gradients on a "correct" backward graph

CatGame's native pi_train LoRA SFT pipeline runs entirely on a patched, in-tree copy of ggml — no PyTorch, no CUDA, no external trainer. That means when something goes wrong in the backward pass, the bug is in our code or in ggml's scheduler, and there is no framework abstraction to blame. This post walks through three coupled investigations that hardened the pipeline: a scheduler CPY race that corrupted LoRA-B gradients, a Vulkan shader rewrite that restored throughput, and a learning-rate sweep that isolated the true divergence cause from the weight-decay scapegoat.

143 → 150 tok/sTiled shader throughput (RTX 3080 Laptop)

1.6402 → 1.5786Train loss, lr=1e-6, 2 epochs (stable)

1.65 → 11.98Train loss, lr=1e-5, 3 epochs (diverges)

276 convsSource-grounded dataset (7 categories)

Bug #1: the mul_mat(cont(transpose(src0))) scheduler race

The original F32 backward path for the LoRA activation gradient computed src0 @ grad via:

ggml_mul_mat(ctx, ggml_cont(ctx, ggml_transpose(ctx, src0)), grad)

Mathematically correct. Operationally fatal. The cont(transpose(src0)) expression creates a CPY node in the backward graph. The ggml scheduler is responsible for materializing that copied/transposed buffer before mul_mat consumes it. On Vulkan, it didn't. The result: mul_mat read uninitialized or stale device-local memory, producing NaNs on CPU and garbage on GPU. The LoRA-A gradient was silently zeroed (because LoRA-B starts at zero, so the chain rule gives dA = B^T @ grad = 0), which masked the bug until LoRA-B began updating.

The fix was not to fix the scheduler. The fix was to stop asking it to materialize a transposed copy of a frozen weight in the backward graph.

The corrected path transposes the gradient instead of the weight, then uses out_prod:

ggml_out_prod(ctx, src0, ggml_to_f32_cont(ctx, ggml_cont(ctx, ggml_transpose(ctx, grad))))

grad is freshly allocated and properly scheduled. src0 is read in its native layout. No CPY race. This restored correctness but exposed the next problem: the only Vulkan out_prod shader in the tree was a naive one-thread-per-output implementation that independently looped over the K dimension.

Bug #2: the throughput cliff

Switching from mul_mat (which has a heavily optimized Vulkan matmul pipeline with cooperative matrix support) to out_prod (which had a naive 256-thread, one-element-per-thread shader) dropped throughput from ~450 tok/s to ~143 tok/s. The shader was correct, but every invocation re-read the entire K dimension from global memory for every output element. On a rank-16 LoRA-B matrix with n=1536, k=16, m=4096, that is 16 redundant global reads per output element across the whole workgroup.

The fix: a 16×16 shared-memory tiled GEMM

The new shader, assets/shaders/pi/out_prod_f32_tiled.comp, is a CatGame-owned 16×16 tiled GEMM. Each workgroup computes a 16×16 tile of the output matrix. The K dimension is processed in chunks of 16. For each chunk, every thread loads one element of A and one of B into shared memory, then all 256 threads accumulate against the shared tile — amortizing global memory reads across the workgroup.

load A tile → shared→load B tile → shared→barrier()→accumulate 16×16→barrier()↻

The shader reads grad as-is (not transposed) and handles the transpose implicitly via swapped indexing, matching the Q4_0 backward convention. It is dispatched via a map_custom1 trigger node in the backward graph — the same pattern used by the Q4_0, Q4_K, and Q6_K quantized backward paths. The trigger is connected to the activation-gradient tensor via a zero-scaled add, forcing the scheduler to run the trigger before the gradient is consumed by the optimizer.

Measured, not estimated: the tiled shader restored correctness and lifted throughput from 143 to 150 tok/s on the RTX 3080 Laptop GPU. The improvement is modest because the F32 LoRA backward is not the dominant bottleneck on this model — the Q4_0/Q6_K quantized weight paths and the forward pass dominate. A cooperative-matrix (coopmat2) variant is the next optimization target for closing the gap to the original ~450 tok/s baseline.

The lr=1e-5 vs lr=1e-6 split

With correctness restored, the next question was why training diverged. Three runs on the 90-conversation / 11-packed-window dataset isolated the cause:

LR	WD	Epoch 1	Epoch 2	Epoch 3	Verdict
1e-5	0.01	1.65	2.44	11.98	Diverges
1e-5	0.001	1.65	2.44	11.98	Diverges (WD not the cause)
1e-5	0.0	1.65	3.61	12.75	Diverges faster (no feedback)
1e-6	0.0	1.6399	1.5782	1.5938	Stable
1e-6	0.01	1.6402	1.5786	1.5937	Stable (user-requested WD)

The split is clean: lr=1e-5 diverges regardless of weight decay. lr=1e-6 is stable with and without wd=0.01. The 90-conversation dataset is too small to support 1e-5 — the gradient signal is noisy and the rank-16 adapter overfits the packed windows, oscillating hard enough to blow past grad_clip=1.0. The user-requested wd=0.01 is safe at 1e-6 and is now the default in config/tools/pi_agent.toml.

Weight decay was not the divergence cause. The learning rate was. This is the kind of false correlation that a controlled sweep catches and a single-run gut check misses.

SSM and NextN backward: already implemented, not skipped

A persistent confusion in the codebase: skip_ssm and skip_nextn sound like they skip the backward pass for SSM/Mamba and Gated DeltaNet layers. They don't. They control LoRA adapter placement — whether LoRA A/B matrices are attached to SSM and NextN weight tensors. The backward passes for GGML_OP_SSM_CONV, GGML_OP_SSM_SCAN, and GGML_OP_GATED_DELTA_NET are implemented in pi/training/ggml/ggml_backward_ssm.cpp and ggml_backward_catgame.cpp, using the same map_custom1 Vulkan dispatch pattern as the out_prod paths. Setting skip_ssm=false would train LoRA on SSM weights (more parameters, slower), not enable a missing backward pass.

TMKL: from post-hoc pullback to a real gradient API

The Target-Masked KL regularizer (from the catastrophic-forgetting prevention literature) removes the target token from both base and adapted distributions, renormalizes, and computes KL on the remaining vocabulary. This preserves the base model's preferences over non-target tokens without opposing the cross-entropy gradient on the target.

The previous implementation was a post-hoc weight pullback: compute TMKL loss, then lerp LoRA weights toward their initial values with strength proportional to (1 + tmkl). This captures the right direction but is not a true gradient. The new compute_tmkl_loss_and_gradient() API computes the exact KL gradient w.r.t. adapted logits:

dKL/dlogits_adapted[v] = p_adapted[v] - p_base_renorm[v] (v != target) dKL/dlogits_adapted[target] = 0

This is the gradient signal that a graph-level TMKL would backpropagate through the LoRA weights. The post-hoc pullback remains the application mechanism (since we cannot add a second forward pass inside the external llama.cpp graph without forking it), but the gradient is now available for a future graph-level integration or for a proper gradient-based correction step.

Dataset expansion: 276 conversations, 7 categories

The 90-conversation dataset was too small for broad conclusions. A source-grounded generator now produces 276 conversations across 7 categories (engine, tools, pi_training, docs, snes, config, shaders) by walking CatGame's own source tree and generating code-explanation, review, architecture, and summary questions for each file. The generator preserves category labels and record boundaries for per-category telemetry. The expanded dataset is at pi_train_large.jsonl on the scratch partition.

80engine

60pi_training

39tools

40docs

30snes

20config

7shaders

What's next

Roadmap

Cooperative-matrix out_prod: a coopmat2 variant of the tiled shader to close the gap to the original ~450 tok/s baseline. The RTX 3080 Laptop GPU advertises NV_coopmat2 support.
Larger-dataset LR sweep: re-test lr=1e-5 on the 276-conversation dataset to see if more data dampens the oscillation that caused divergence on the 90-conversation set.
Graph-level TMKL integration: wire compute_tmkl_loss_and_gradient() into a proper gradient correction step, replacing the L2 pullback approximation.
SSM/NextN LoRA training: run with skip_ssm=false skip_nextn=false to measure the parameter-count and throughput impact of training LoRA on hybrid model layers.

Caveats and limitations

Throughput numbers are measured on one GPU (RTX 3080 Laptop GPU, 8 GB VRAM, Vulkan backend). Results will differ on other hardware. The 150 tok/s figure is the measured throughput with the tiled shader on the 90-conversation dataset at 4096 context, batch 1, rank 16.
The 143 → 150 tok/s improvement is modest. The F32 LoRA backward is not the dominant bottleneck on this model; the quantized weight paths and forward pass dominate. A coopmat2 shader is needed for a larger gain.
The lr=1e-5 divergence conclusion is specific to the 90-conversation dataset. The 276-conversation dataset may support a higher LR. This has not yet been tested.
TMKL gradient API is implemented but not yet wired into the training loop as a gradient correction. The post-hoc pullback remains the active mechanism.
The dataset generator produces synthetic Q&A from source file structure (includes, declarations, comments). The responses are templated, not human-written. This is sufficient for testing training stability but not for producing a production-quality adapter.
SSM/NextN backward is implemented but not yet trained with skip_ssm=false. Correctness of the backward pass on real training data with LoRA attached to SSM weights is unverified.