Three root-cause bugs made the collapse-recovery distillation mechanism
a silent no-op. Diagnosed via 50-epoch smoke test trajectory: Q-gap
peaked at 1.2 in epoch 2 then collapsed irreversibly by epoch 6, with
HEALTH_DIAG reporting `distill=off` throughout despite the trigger
conditions firing every epoch. Fixed each bug in turn and re-ran:
Q-gap now stabilizes above 0.18 from epoch 33 onwards, final 1.18.
Bug 1: timing
apply_distillation_gradient ran at epoch-boundary and SAXPY'd into
grad_buf. But the next training step's graph_forward.replay() starts
with `cuMemsetD32Async(grad_buf, 0, total_params)` — wiping the
contribution before any Adam update could see it. Moved SAXPY into
the per-step aux-op phase (between graph_forward and graph_adam),
matching the cadence of CQL/IQN/ensemble gradients.
Bug 2: CUDA Graph scalar baking
Moving SAXPY to per-step hit a deeper issue: graph capture bakes
kernel scalar args at capture time. The alpha value was captured
at 0.0 (initial) and never updated across replays, regardless of
per-epoch recomputation. Fix: new dqn_distill_saxpy_kernel reads
health from isv_signals[LEARNING_HEALTH_INDEX=12] directly and
computes alpha in-kernel. `distill_best_buf` is a stable device
pointer; its contents are DtoD-refreshed at epoch boundary when
maybe_snapshot_params accepts a new best. Zero CPU writes on any
path — pure GPU dataflow.
Bug 3: snapshot gate using wrong signal
The snapshot gate passed `self.last_q_gap` (an EMA that was stuck
at 0 due to broken propagation — see companion commit). Gate was
`health ≥ 0.65 OR winrate_fallback`, neither of which opened in
runs where high-q_gap epochs and high-winrate epochs don't overlap.
Replaced with q_gap-primary gate: `epoch_q_gap ≥ dynamic_floor`,
where the floor is `0.5 × decaying_peak` scaled by a per-epoch 0.99
decay. Adapts to network size automatically (production peaks at
~0.5 → floor 0.25; smoke test peaks at ~0.05 → floor 0.025).
Also drops the winrate-fallback "inflate health to 0.75" hack from
training_loop — q_gap is the direct measure of what distillation
preserves, no proxies needed.
Verified: local E1 smoke test (RTX 3050 Ti, 50 epochs) shows
distillation engaging from epoch 2 onwards and keeping Q-gap above
0.18 for epochs 33-50. Production L40S 50-epoch run pending deploy.
Files changed:
- dqn_utility_kernels.cu: new dqn_distill_saxpy_kernel (numerically
unchanged from saxpy_f32_kernel; alpha computed from ISV per-thread)
- gpu_dqn_trainer.rs: distill_saxpy_aux kernel handle, distill_best_buf
stable device buffer initialized from params at construction,
apply_distillation_gradient() rewritten, mirror_best_snapshot_to_distill_buf()
invoked on snapshot acceptance, maybe_snapshot_params uses dynamic
q_gap floor via SnapshotRing::observe_q_gap/dynamic_q_gap_floor
- q_snapshot.rs: SnapshotRing grows max_q_gap_observed decaying-peak
tracker, observe_q_gap() + dynamic_q_gap_floor() helpers,
MIN_SNAPSHOT_Q_GAP constant dropped in favor of relative floor,
module docstring rewritten to explain q_gap-primary gate
- fused_training.rs: set_distill_alpha + distill_alpha_per_step field
dropped (no longer needed — kernel reads ISV directly),
submit_aux_ops calls apply_distillation_gradient() unconditionally
- training_loop.rs: snapshot call simplified — passes epoch_q_gap
(raw, not the stuck EMA), drops winrate fallback and distill alpha
plumbing; also calls fused.update_eval_v_range() in the epoch-end
Q-stats block (the path previously writing to per_branch_q_gap_ema
was disabled via `if false` guard, leaving the health EMA frozen
at zero)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
24f96ab78b