The Flat Loss Mystery: LoRA Training Diagnostics on Vulkan

We ran four controlled experiments on our native Vulkan LoRA training pipeline, varying KL regularization, weight decay, learning rate, and rsLoRA scale. Every single run produced the same epoch-level loss to four decimal places. Here's what we found — and what it means for how you measure LoRA training progress.

The Setup

Our pi_train tool is a from-scratch LoRA SFT trainer built on top of llama.cpp's ggml optimizer. It runs entirely on Vulkan — no CUDA, no PyTorch, no Hugging Face Trainer. The target model is Qwen3.5-2B (F16), fine-tuned on 3,079 CatGame code and documentation conversations packed into 4096-token windows with 99.6% utilization.

The training pipeline includes:

The Mystery

When we started tuning KL regularization, we noticed something strange. The epoch-level training loss was identical across all KL settings — 0.0, 0.01, 0.03, and 0.1 all produced train_loss=2.0829 for epoch 1, 2.0879 for epoch 2, and 2.0855 for epoch 3. To four decimal places. Every time.

The Experiments

We designed a controlled experiment matrix using the scientific method — changing one variable at a time while keeping everything else fixed:

Plus three prior runs with different KL and WD values at LR=1e-4 — all producing the same numbers. Seven experiments total, all identical.

What We Found

KL Regularization: Three Approaches Tested

1. L2 Weight Pullback (Original)

Our first KL implementation was a post-hoc weight pullback: after each optimizer step, pull LoRA weights slightly toward their initial values.

theta -= lr * lambda * (theta - theta_0)
// Equivalent to adding 0.5 * lambda * mean(||theta - theta_0||^2) to loss

The problem: with lr=1e-4 and lambda=0.03, the pullback step is 3e-6 — moving weights by 0.0003% per step. Effectively zero. This is why all KL values produced identical results.

2. True Graph-Level KL (Attempted)

We tried injecting the KL gradient directly into ggml_opt_grad_acc() accumulators before the optimizer step. This would make the optimizer see CE_grad + KL_grad and apply momentum/clipping to the combined gradient.

3. Target-Masked KL (TMKL)

Based on research from Mask the Target (arXiv:2605.29498), we implemented TMKL — the correct KL approach for LoRA SFT:

// At each supervised position:
// 1. Get base model distribution p_base (LoRA scale=0)
// 2. Get adapted model distribution p_adapted (LoRA scale=1)
// 3. Remove target token y from both, renormalize
// 4. KL(adapted || base) over non-target vocabulary
//
// L = CE + lambda * TMKL

TMKL preserves the base model's preferences over non-target tokens without opposing the cross-entropy gradient. Standard KL (and our L2 pullback) directly fights CE — the base model assigns low probability to target-domain tokens that CE must learn, so any KL pull toward the base model opposes learning.

rsLoRA Scale Analysis

rsLoRA (rank-stabilized LoRA) adjusts the scaling factor from alpha/rank to alpha/sqrt(rank). With our defaults:

// Standard LoRA:  scale = alpha / rank = 128 / 64 = 2.0
// rsLoRA (alpha=2r): scale = 1024 / 64 = 16.0  (8x amplification!)
// rsLoRA (alpha=r):  scale = 64 / 64 = 1.0     (via stored alpha=64*sqrt(64)=512)

We tested all three scales (16.0, 8.0, 2.0). The loss was identical in all cases — but this doesn't mean scale doesn't matter. It means the loss metric can't distinguish between them. The actual weight updates and per-category effects may differ, but the cumulative average washes them out.

What This Teaches Us

Lessons for LoRA Practitioners

• Check weight checksums — if MD5s differ between init and trained adapters, the LoRA IS learning, even if the loss looks flat.
• Track per-category losses — aggregate averages hide real progress. Some categories improve while others regress, netting zero.
• KL pullback step size matters — lr * lambda must be large enough to have an effect. At 3e-6, it's invisible.
• Weight decay on LoRA IS L2-SP regularization — it pulls toward W_init, not zero. Having both WD and KL is redundant double-regularization.
• TMKL is the correct KL for SFT — standard KL opposes CE under distribution shift. Target-masking removes the conflict.
• Validation sets need to be large enough — 28 data points can't show subtle per-category shifts.

The Vulkan Training Pipeline

All experiments ran on our native Vulkan compute backend — no CUDA dependency. The pipeline achieves 850 tok/s on an 8GB GPU with 228 LoRA tensors across all model layers (attention, FFN, SSM, NextN). This is a from-scratch implementation: the forward pass, backward pass, gradient accumulation, optimizer, and telemetry are all CatGame-owned code, with only the base ggml graph engine as an external dependency.

Next Steps

• Add per-batch loss delta metrics that show change vs first epoch
• Investigate whether validation forward pass applies LoRA adapter
• Try much higher learning rates (1e-3, 5e-3) to see if loss moves
• Implement true graph-level TMKL with a custom forward pass path
• Use held-out perplexity as the primary training metric

Caveats and Limitations

• All experiments used the same fixed seed (42) and dataset. The identical losses across hyperparameters may partly reflect dataset determinism rather than pure metric insensitivity.
• The TMKL implementation is infrastructure-only — the post-hoc pullback approximation doesn't capture the target-masking benefit. Full TMKL requires graph-level integration that is a roadmap item.
• The validation set (28 data points, 307 conversations) is small. The frozen val_loss may be a sample-size artifact rather than a LoRA application bug.
• Per-category loss changes are computed from different data point assignments across epochs (due to reshuffling), so direct per-category comparison has noise.
• These results are specific to our Vulkan + ggml + Qwen3.5-2B + CatGame dataset combination. Other model architectures, datasets, or backends may show different behavior.