The Silent Killer: Unmasked Padding in LoRA Cross-Entropy Loss

The Symptom

We were training a LoRA adapter on Qwen3.5-2B — a hybrid model with 25 layers interleaving softmax attention and Gated DeltaNet (a state-space / linear attention variant). The architecture matters because it changes which layers get LoRA adapters and how gradients flow. But the bug we found wasn't in the architecture. It was in something far more mundane.

The training loss looked fine for the first epoch. Then it started climbing. By epoch 3, validation loss had gone from 3.64 to 4.47 — the model was getting worse with more training. The LoRA adapter drift was growing exponentially: 3.37 → 6.66 → 9.79. And the last batch of each epoch showed a bizarre loss spike above 7, sometimes 9.

The Investigation

We followed the scientific method: form a hypothesis, change one variable, measure, compare. The leading suspects were learning rate, gradient clipping, hybrid SSM gradients, KL regularization, and dataset packing. One by one, we tested them. None explained the divergence pattern.

Then we traced the loss computation path. The training loop uses ggml_cross_entropy_loss — a fused op that computes both forward loss and backward gradients in a single pass for numerical stability. The external implementation in ops.cpp revealed the truth:

// Forward: loss = -sum(log_softmax(logits) * labels) / nr
// Backward: grad = (softmax - labels) / nr * grad_loss
//
// nr = total number of rows (ALL positions, including padding)

The op divides by nr — the total number of positions, including padding. And the global g_pad_token_id? It was only used for accuracy accounting (counting correct predictions), not for loss masking. Padding positions had one-hot labels at the pad token index, meaning the model was being trained to predict pad_token at every padding position.

The Fix

The fix required two changes, both in CatGame-owned code (no external modifications):

1. Zero out pad_token labels

Before graph compute, read the pad_token row from the labels tensor. For each position where the pad_token label is 1.0 (indicating padding), zero it out. This removes the spurious "learn to predict pad_token" gradient signal. The remaining gradient for padding positions is softmax / n_valid — a benign uniform push.

2. Loss scale correction

The fused op divides by nr (total positions). After masking, we need division by n_valid (valid positions only). A runtime scalar input loss_scale = n_total / n_valid is multiplied with the CE loss in the graph:

// Graph: loss = ce_loss * loss_scale
//   ce_loss = -sum(log_softmax * labels_zeroed) / n_total
//   loss_scale = n_total / n_valid
//   => loss = -sum(log_softmax * labels_zeroed) / n_valid
//
// Backward: grad = (softmax - labels_zeroed) / n_valid
//   Valid positions:  grad = (softmax - one_hot) / n_valid  (CORRECT)
//   Padding positions: grad = softmax / n_valid              (benign)

The Results

We ran a controlled experiment: same dataset (276 conversations), same hyperparameters (lr=1e-6, rank=16, alpha=32, 3 epochs, skip_ssm=true), with and without the masked CE fix.

The divergence was completely eliminated. Validation loss stayed flat at 1.11 across all three epochs. LoRA drift grew linearly (0.34 → 0.97) instead of exponentially (3.37 → 9.79). The last-batch loss spike disappeared entirely. And training throughput was unchanged at 159 tokens/second — the masking overhead is negligible.

Why It Was So Subtle

Three things made this bug hard to find:

• Misleading comments. The codebase had comments saying padding labels were "harmless" because the model would "learn to predict pad_token, which is harmless." This is wrong for SFT — you don't want the model wasting capacity on pad prediction.
• The validation set was 52% padding. The val_loss=4.44 looked like a real signal of poor model quality. It was actually dominated by padding positions being penalized for not predicting pad_token. After masking, the true val_loss was 0.83 — the model was already good.
• The fused op looked correct. The cross-entropy math was numerically stable and correct for the positions it computed over. The bug was in which positions it computed over, not in how it computed.

Hybrid Architecture Lessons

Separate research confirmed that our skip_ssm=true setting is correct for Qwen3.5. The paper "Where Should LoRA Go?" (arXiv 2604.22127) studied Qwen3.5-0.8B and found that adapting the recurrent backbone is destructive in sequential hybrids (−14.8 pp on GSM8K), while attention-only LoRA outperforms full-model adaptation with 5–10× fewer trainable parameters.

The SSM/DeltaNet weights govern hidden state transitions. Even minor scale drift in these tensors destabilizes the recurrence, causing context collapse and repetitive looping at long sequences. Keeping LoRA adapters off the SSM layers avoids this failure mode entirely.

LR Sweep

With the masked CE fix in place, we swept the learning rate to find the new stable range:

The masked CE fix expanded the stable LR range. Before the fix, even lr=1e-5 diverged catastrophically. After the fix, lr=5e-6 shows only mild overfitting — a 5× improvement in LR headroom.

Caveats

• Padding gradient is not zero. The fix zeros the pad_token label, but padding positions still contribute softmax / n_valid gradient (a benign uniform push). A truly masked op would zero these entirely, but the current approach is numerically correct for valid positions.
• Eval improvement is small. With lr=1e-6, the LoRA improves eval loss by only −0.06%. The base model already performs well on this data (it was used to generate the training set). More training data (3,079 records in the balanced dataset) should improve generalization.
• Host-side masking overhead. The current implementation reads the pad_token row position-by-position via ggml_backend_tensor_get. For 128 tokens per batch, this is 256 backend tensor operations. The overhead is negligible (159 tok/s unchanged), but could be optimized with a bulk read.
• Architecture-specific findings. The skip_ssm=true recommendation is specific to sequential hybrids like Qwen3.5. Parallel hybrids (e.g., Falcon-H1) may benefit from SSM adaptation.

Files Changed

• pi/training/ggml/ggml_opt_internal.h — Added loss_scale tensor to opt context
• pi/training/ggml/ggml_opt_build.cpp — Added loss_scale runtime input, graph multiplies CE loss by it
• pi/training/ggml/ggml_opt_eval.cpp — Before compute: zero pad labels, count valid, set loss_scale = n_total / n_valid
• config/tools/pi_agent.toml — Updated training defaults with masked CE findings