diag+docs(dqn): trunk-slice grad decomposition + stale IQN-trunk doc fix

Extends grad_decomp_kernel to snapshot the trunk tensor slice (tensors
0..4 = w_s1, b_s1, w_s2, b_s2) in addition to the existing direction +
magnitude branch slices (8..12 / 12..16). Adds a new HEALTH_DIAG group:

  grad_trunk [iqn=<abs> ens=<abs> c51=<abs> cql=<abs> distill=<abs>
              rec=<abs> pred=<abs> cql_sx=<abs> c51_bs=<abs>]

Prior grad_split_bwd / grad_split_aux groups report mag_norm / dir_norm
ratios per loss component, computed over branch-head tensors only. That
measurement range structurally reports 0.0000 for any loss component
that writes exclusively to the trunk — IQN and Ens in particular. This
caused the persistent misdiagnosis that IQN-to-trunk was not wired; the
prior scoping in /tmp/foxhunt_research/iqn-to-trunk-wiring-scoping.md
confirmed the wiring is live (apply_iqn_trunk_gradient at
gpu_dqn_trainer.rs:4882) and that the zero reading was a blind spot in
the measurement pipeline.

Smoke confirms the diagnostic: after iqn_readiness ramps up (late
epochs), grad_trunk reports iqn=100..381 (real trunk SAXPY amplitude),
ens=0.07..3.57, c51=2.46..8.91 (value-head dueling path contributes
through trunk), while cql/cql_sx/distill/rec/pred stay near-zero — a
clean diagnostic baseline.

Also fixes stale documentation at dual-distributional-c51-iqn-design.md
that claimed "IQN trains in isolation — its gradients don't flow back
to the shared trunk": reworded to reflect current wired state with
file:function citation and explicit iqn_readiness gating note. Updated
the "What Changes" table ("IQN training") and "Implementation Order"
(Phase 1 marked DONE) with the same citation.

Changes:
  - grad_decomp_kernel.cu: per-component result slot 2 → 3 floats
    (mag_norm, dir_norm, trunk_norm); extra __shared__ sum_trunk +
    tree reduction; new grad_trunk_start/trunk_len kernel args.
  - gpu_dqn_trainer.rs: pinned result buffer 18 → 27 floats; snapshot
    now does two copy_f32 passes (trunk → dst[0..trunk_len), branch →
    dst[trunk_len..]); per-component slot offsets 0/2/… → 0/3/…;
    grad_component_norms_trunk cached field + accessor; compute trunk
    range from padded_byte_offset(&param_sizes, 0..4).
  - fused_training.rs: grad_trunk_norms_by_component() + per-component
    grad_trunk_*_abs() accessors.
  - training_loop.rs: HEALTH_DIAG emits new grad_trunk group ordered
    [iqn ens c51 cql distill rec pred cql_sx c51_bs]; extended doc
    comment explaining the three groups' roles.
  - design spec (Problem #1 + What Changes row + Implementation Order):
    stale "IQN trains in isolation" replaced by current wired-state
    description, cites gpu_dqn_trainer.rs:4882 and readiness ramp at
    gpu_dqn_trainer.rs:4228-4243.

Pure diagnostic — no training dynamics change, no atomicAdd, no tuning
knobs, no TF32 changes. Smokes unaffected (magnitude_distribution H10
regression pre-exists on HEAD 810b3c570; 4 other smokes pass).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-22 22:03:20 +02:00
parent 810b3c5703
commit b8cd4e1d94
5 changed files with 300 additions and 113 deletions

View File

@@ -3,7 +3,17 @@
## Problem
The system has both C51 (101 atoms) and IQN (32 quantiles) implemented, but:
1. IQN trains in isolation — its gradients don't flow back to the shared trunk
1. **[Phase 1 DONE]** IQN's backward gradient is now wired into the shared
trunk via `apply_iqn_trunk_gradient` (`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs:4882-5034`),
captured unconditionally inside the aux child graph from
`crates/ml/src/trainers/dqn/fused_training.rs:1582-1587`. The effective
scale is `iqn_lambda * iqn_readiness * iqn_budget`, where
`iqn_readiness` is a loss-improvement-driven ramp (0→1) at
`gpu_dqn_trainer.rs:4228-4243`. Trunk contribution is measurable via
the Task 2.Z `grad_trunk [iqn=...]` HEALTH_DIAG group — the previous
`grad_split_bwd [iqn=0.0000]` reading was a measurement blind spot in
the branch-only grad_decomp snapshot range, not evidence of missing
wiring.
2. IQN's quantile output is never used for action selection or risk management
3. `use_cvar_action_selection` exists in config but is hardcoded to `false`
4. The trunk learns only from C51's cross-entropy loss, missing the risk signal from IQN
@@ -56,7 +66,7 @@ Wire the existing IQN dual-head for two purposes:
| Component | Today | After | Effort |
|-----------|-------|-------|--------|
| C51 training | Working | No change | 0 |
| IQN training | Working, disconnected | Gradient flows to trunk | LOW |
| IQN training | Working, trunk gradient wired (Phase 1 landed — `apply_iqn_trunk_gradient`) | No further change | DONE |
| IQN inference | Kernel compiled, never called | Wired for CVaR | MEDIUM |
| CVaR computation | Not implemented | New trivial CUDA kernel | LOW |
| C51/IQN disagreement | Not implemented | |E_c51 - E_iqn| | LOW |
@@ -67,7 +77,11 @@ Wire the existing IQN dual-head for two purposes:
## Implementation Order
1. Wire IQN gradient to shared trunk (Phase 1)
1. ~~Wire IQN gradient to shared trunk (Phase 1)~~**DONE**
(`apply_iqn_trunk_gradient` at `gpu_dqn_trainer.rs:4882-5034`,
scheduled from `fused_training.rs:1582-1587`, scale
`iqn_lambda * iqn_readiness * iqn_budget`). Measurable via the
Task 2.Z `grad_trunk [iqn=...]` HEALTH_DIAG group.
2. Wire IQN inference + CVaR computation (Phase 2)
3. Add disagreement uncertainty signal (Phase 3)
4. Hyperopt integration for cvar_alpha (Phase 4)