GRPO: ValidTokenCount reads 0 for a full step while Losses accumulate — per-micro vtc semantics + rare last-micro corruption

[flaviusburca]

Symptom

Rare (3 occurrences over ~150 GRPO steps across two engine builds): a training step logs vtc=0 with hundreds of micro-batches, producing garbage step metrics (masked 2384-3052%, kl 35-42, loss 42-59) and exploded pre-clip grad norms (0.77-0.96). The zero-lr guard from #72 neutralizes it (policy untouched), but the root cause is open.

Established by investigation (2026-08-04)

1. ValidTokenCount is per-micro in the GRPO path: forward_for_grpo unconditionally zeroes Losses/ValidTokenCount/CorrectCount every micro (dsl_model_execution.cpp:795). The trainer's step-end get_valid_token_count(0) therefore returns the last micro's count, not a step total. Verified: a healthy step logged vtc=136 and its last micro's loss-mask count is exactly 136 (decoded from the transport batch).
2. On corrupted steps the CE kernels did run and write Losses (logged loss ≈ raw -logprob sum over one micro's ~125-170 valid tokens), while vtc read 0 — same kernels write both (fused_lm_head_loss.cpp:417-441), so a stuck kv_sweep skip (fused_lm_head_loss.cpp:242) is ruled out (it would zero both).
3. Data ruled out: both preserved batches decode cleanly; no micro has a zero loss-mask; last micros carry 125-167 valid tokens (rollouts_step73_FORENSIC, rollouts_step82_FORENSIC in the pool-D output dir, plus the cap3-era step-14 batch).
4. dispatch_pp_apply_optimizer's host overwrite clamps to ≥1 — ruled out.

Remaining candidates

• The last micro's chunked_cross_entropy_forward/fused_cross_entropy_forward contributing 0 to the vtc atomic while writing losses (kernel-internal accumulation edge, possibly interacting with lmhead_chunks nano-slicing or the chunked phase-A/B schedule).
• A post-backward re-zeroing of ValidTokenCount racing the step-end read (any path that calls fill_zero on run-state loss buffers between the last backward_grpo and the read).

Next step

Deterministic replay: feed a forensic batch through a single-GPU debug trainer (same config: fp8-hybrid, sequence_chunks 4, lmhead_chunks 12, single-sample bins) with SUROGATE_CHUNK_TRACE=1 and a vtc D2H readback after each micro — bisects kernel-edge vs post-hoc-zeroing in one run. Batches and configs preserved; replay harness is the open work.

Also worth considering

Step-end consumers treating vtc as a step total (metric normalization; global_norm_sqrt token scale) may deserve an explicit step-accumulated counter regardless of this bug — the per-micro semantics make every step-level metric a last-micro sample.

🤖 Generated with Claude Code

[flaviusburca]

Root cause identified (no replay needed)

forward_for_grpo zeroes Losses/ValidTokenCount unconditionally per invocation (dsl_model_execution.cpp:795). Chunked GRPO invokes it per chunk, and phase B executes chunks in reverse — the step's final invocation is always the last micro's chunk 0, so the step-end read returns chunk 0's valid count of the last sample.

Verified predictions (decoded transport batches):

• healthy steps: last micro single-chunk, mask counts 136/135 == logged vtc 136/135 (exact)
• corrupted steps 73/82: last micro is a 2-chunk sample with valid-per-chunk [0, 125] / [0, 167] — the prompt fills chunk 0 completely → vtc reads 0

Incidence arithmetic: P(>1536-token-prompt sample) × P(packed last) ≈ few % × 1/330 ≈ once per 50-100 steps — matches 3 in ~150 observed. Deterministic, data-order-dependent, engine-version-independent.

Consequences

1. The vtc=0 crash class (this issue): guarded by fix(grpo): never apply an optimizer step whose token accounting is corrupt #72's zero-lr flush; root fix = accumulate a step-level valid-token counter (zero once at micro 0, accumulate across chunk invocations) for the step-end consumers, or key the per-chunk zeroing so reverse-order phase B doesn't leave chunk-0-only state.
2. Broader: in chunked GRPO, every step-level consumer of ValidTokenCount (metric normalization, the token-scale in the grad-norm path when mUseTokenScale) has been sampling one chunk's count, not the step total — consistent across steps (so absorbed into effective scaling) but semantically wrong and the direct cause of the metric garbage / norm explosion in the zero case.

Cheap interim (python-side, no engine rebuild): the trainer can compute the true step vtc as Σ loss_mask over the batch it already holds and use that for metrics/guard decisions; the engine counter then only feeds the norm path until the engine fix lands.

🤖 Generated with Claude Code

[flaviusburca]

Follow-up: the non-zero form distorts the grad-norm instrument campaign-wide

Regressing 76 logged steps: corr(reported grad_norm, 1/vtc) = 0.83 — the reported norm is dominated by the inverse of the last micro's chunk-0 valid count (the token-scale denominator), not by gradient magnitude. Largest observed artifacts: vtc=38 → norm 0.088; vtc=30 → norm 0.264 on a batch whose loss/kl/mask were fully normal.

Consequence beyond the vtc=0 crash: any grad-clip threshold reacts to this phantom norm and suppresses real updates on low-vtc steps (0.05 clip × norm 0.26 ⇒ applied update cut ~5x), and norm-based health rules fire on packing noise.

Interim mitigation (python-side, gradient-neutral, shipped in our campaign trainer): reorder each batch so the micro with the maximum chunk-0 valid count is last — pins the denominator at the batch max, restores cross-step comparability, and eliminates the 0 case. The proper engine fix (step-accumulated counter) supersedes this.

🤖 Generated with Claude Code

[flaviusburca]

Boundary case: the packing-order mitigation cannot cover pure-long-prompt batches

Live occurrence #4 (step 128 of the campaign): a batch of 298 micros where every sample's prompt exceeds the 1536-token chunk size (near-pure office/repo composition — those lanes carry 2-4k-token contexts). max(chunk0_valid) == 0 across the whole batch, so no reordering can produce a nonzero last-micro chunk-0 count and the engine necessarily reads vtc=0. The #72 zero-lr guard discarded the step safely (policy untouched).

This bounds what trainer-side mitigation can achieve: the reorder guard (PR #75) fixes mixed batches — the overwhelming majority — and stabilizes the norm denominator, but batches whose lane composition is uniformly long-prompt (measured ~3% here) are structurally un-fixable from the data side. They are currently paid as one discarded batch each.

Sharpens the engine-fix requirement: the step-accumulated ValidTokenCount counter is the only complete solution.

🤖 Generated with Claude Code