POC: stop-gradient through context tokens (row-sparse backbone backward)

[samsja]

POC on top of #3326 — requires that PR's keep-mask plumbing. Changes the training gradient by design; not a pure optimization.

What

model.stop_grad_context_tokens (default false): keep the forward pass bit-identical (context K/V are still built with current weights), but make the backbone backward row-sparse — gradients only propagate through tokens some loss component reads.

Mechanism (layers/row_sparse.py):

• mask_grad at every decoder-layer input: identity forward, backward zeroes context-row grads. This stops gradient propagation into context hidden states and everything below them.
• row_sparse_linear for q/o and the dense MLP: full-row forward, backward compacted to kept rows (dW = dY[keep]^T @ X[keep], dX[keep] = dY[keep] @ W). Valid because under the mask_grad cut their upstream grads are exactly zero on context rows.
• k/v projections keep a full backward: kept queries attend to context k/v, so dK/dV on context rows are nonzero and feed the k/v weight grads — the model keeps learning how to read context; it stops learning how to represent it.

Reuses the label-aligned keep_mask from #3326 (hence the config validator requiring skip_masked_lm_head_tokens + fused head). Wired for the custom Qwen3 dense backbone only; other models raise.

Measured (4x RTX 3090, Qwen3-1.7B)

Single-GPU fwd+bwd microbench, seq 2048, on top of head-skip:

kept fraction	head-skip only	+ stop-grad	speedup
0.62	499.8 ms	448.4 ms	1.11x
0.50	490.0 ms	422.9 ms	1.16x
0.20	453.3 ms	336.7 ms	1.35x
0.10	440.5 ms	302.4 ms	1.46x

SFT A/B (wordle, 12 steps, 4 GPUs, micro_batch_size=2, kept fraction 0.62): loss curves track closely (step 12: 0.4719 off vs 0.4816 on; grad norms lower with stop-grad, as expected with fewer gradient paths). No e2e speedup visible in this config — the runs were data-loader-bound (~40.8s fwd+bwd wall vs ~15s of GPU compute at these shapes; tokenization/packing dominates on this box), which masks the compute win the microbench isolates.

Correctness tests (tests/unit/train/test_row_sparse.py, all passing on GPU):

• row-sparse linear backward == dense torch.where(detach) reference (grads exact)
• all-masked batch keeps a dummy row so weight grads/collectives survive
• full Qwen3 path: forward identical with/without the flag; context-only tokens get exactly zero embedding grad

Known gaps (POC)

• Qwen3 dense only; MoE/VLM/other backbones raise. LoRA/quantized linears fall back to dense backward.
• Attention-kernel backward itself still runs full-length (FA2 computes dQ/dK/dV for all rows); only the projection/MLP GEMMs are row-sparse. A query-split attention (grad for kept queries only) is the next ~15-20%.
• backbone_keep_index uses nonzero() (one D2H sync per micro-batch, same as the head-skip path in Don't compute the logprobs and their gradients for tokens no loss com… #3326).
• Interaction with activation checkpointing is unvalidated (recompute is full-length; savings shrink).
• Needs a real convergence study before anything beyond experiments — this changes what the model learns.