gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

spec(sp7): loss-balance controller — adapt CQL/C51 budgets to grad-norm ratio

Browse Source 
Outcome-driven controller replacing Pearl 2's hardcoded-0 CQL budget and
floor-pinned C51 budget with multiplicative ratio adaptation. Targets
`grad_split_bwd cql/iqn = 2.0` and `c51/iqn = 1.0` per slice (trunk, dir,
mag). Slow EMA (α=0.01) consistent with existing controller pearls; Pearl
A sentinel-bootstrap on fold boundary; Pearl D Wiener-optimal smoothing.

Replaces ghost docstring at fused_training.rs:3409 ("0.10×(1−regime)×health"
formula was never implemented). Pearl 2 keeps owning IQN budget (reference
denominator) and FLATNESS_BASE (consumed by NoisyNet σ pearl); stops
writing CQL/C51/ENS slots so the new controller is the sole driver.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-03 00:03:14 +02:00
parent 04df20bdc3
commit 0808ebe3e5
1 changed files with 250 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

250
docs/superpowers/specs/2026-05-03-sp7-loss-balance-controller-design.md Normal file
View File

@@ -0,0 +1,250 @@
# SP7: Loss-Balance Controller — Design
**Status**: approved (2026-05-03)
**Branch**: `sp5-magnitude-differentiation` (extending SP5 closure)
**Predecessor**: SP5 close-out (commits `ab378327e` … `04df20bdc`)
## Goal
Replace the dormant per-branch CQL budget output (hardcoded to 0 by Pearl 2, floored to 0.02 in the consumer) and the floor-pinned per-branch C51 budget output with a single **outcome-driven controller** that adapts both heads' per-branch budgets so each loss head's gradient contribution stays in a target ratio relative to IQN.
This unblocks magnitude differentiation: at production scale (50-epoch baseline at HEAD `2fb7d7f57`, smoke `train-multi-seed-n4qv2` ep4-7) we observed `grad_split_bwd cql=4.13` vs `iqn=0.07` (≈60×), with magnitude Q-values within 0.003 of each other, producing `eval_dist eq=1.000` (full eval-collapse to Quarter via argmax).
## Non-goals
- Not redesigning Pearl 2. Pearl 2 keeps writing `flatness[b]` to `FLATNESS_BASE` (consumed by NoisyNet σ — `pearl_per_branch_noisy_sigma`) and `budget_iqn[b]` to `BUDGET_IQN_BASE`. IQN is the reference; needs to stay stable.
- Not adding new diagnostic instrumentation. The grad-norm signals the controller consumes already exist (per-loss `grad_decomp_launch_*` pinned slots, populated every step).
- Not adapting trail/distill/rec/pred budgets. Out of scope. Same controller pattern can be applied later.
## Architecture
```
Existing per-step pipeline (every training step)
────────────────────────────────────────────────
Forward pass → Backward (IQN, CQL, C51, Ens, distill, …)
│
Each loss component runs `grad_decomp_launch_*`
which writes a 3-float pinned slot:
[mag_norm, dir_norm, trunk_norm]
│
▼
┌────────────────────────────────────────────────────────┐
│ loss_balance_controller_kernel (NEW — single block, 8 threads)│
│ │
│ Reads: │
│ iqn_decomp_pinned[3] ← grad_decomp_launch_iqn │
│ cql_decomp_pinned[3] ← grad_decomp_launch_cql_sx │
│ c51_decomp_pinned[3] ← grad_decomp_launch_c51 │
│ isv[BUDGET_CQL_BASE+0..4] ← previous step │
│ isv[BUDGET_C51_BASE+0..4] ← previous step │
│ │
│ Computes (per head h ∈ {CQL, C51}, per slice s): │
│ actual_ratio = h_norm[s] / max(iqn_norm[s], EPS) │
│ correction = target_ratio[h] / max(ratio, EPS) │
│ new_budget = (1-α)·old + α·(old × correction) │
│ new_budget = clamp(new_budget, EPS, MAX_BUDGET) │
│ │
│ Writes (raw, pre-smoothing): │
│ scratch[BUDGET_CQL_BASE+0..4] │
│ scratch[BUDGET_C51_BASE+0..4] │
└────────────────────────────────────────────────────────┘
│
▼
apply_pearls_ad_kernel (existing) — Pearl A sentinel-bootstrap
+ Pearl D smoothing → ISV[BUDGET_CQL_BASE+0..4],
ISV[BUDGET_C51_BASE+0..4]
│
▼
Next step's compute_adaptive_budgets() reads these slots
→ drives apply_cql_saxpy / apply_c51_budget_scale.
```
## Inputs (existing — no new producers needed)
All three are pinned 3-float slots `[mag_norm, dir_norm, trunk_norm]`, computed every step by the corresponding `grad_decomp_launch_*` after each loss component's backward pass. The pinned device pointers are already accessible as kernel args.
| Field | Source | Wired? |
|---|---|---|
| `iqn_norm[mag, dir, trunk]` | `grad_decomp_launch_iqn` (per-step) | yes |
| `cql_norm[mag, dir, trunk]` | `grad_decomp_launch_cql_sx` (per-step) | yes |
| `c51_norm[mag, dir, trunk]` | `grad_decomp_launch_c51` (per-step) | yes |
The `cql_sx` decomp captures the SAXPY-applied delta (post-budget); `iqn` and `c51` capture the pre-budget delta. We use `cql_sx` for parity with what actually lands in `grad_buf`.
## Outputs (existing dormant slots)
| Slot | Current state | New state |
|---|---|---|
| `BUDGET_CQL_BASE+0..4` (ISV 198..202) | Pearl 2 writes 0 → floors to 0.02 in `compute_adaptive_budgets` | Controller writes adaptive per-branch budget |
| `BUDGET_C51_BASE+0..4` (ISV 190..194) | Pearl 2 writes `flatness × …` → in production drops below floor 0.05 | Controller writes adaptive per-branch budget |
`BUDGET_ENS_BASE+0..4` and `BUDGET_IQN_BASE+0..4` remain Pearl 2's responsibility for now (Ens is residual; IQN is reference).
## Math
For each loss head `h ∈ {CQL, C51}` and each slice `s ∈ {dir, mag, trunk}`:
```
target_ratio[CQL] = TARGET_CQL_OVER_IQN // Invariant 1: 2.0
target_ratio[C51] = TARGET_C51_OVER_IQN // Invariant 1: 1.0
iqn_n = max(iqn_norm[s], EPS_DIV)
h_n = h_norm[s]
actual_ratio = h_n / iqn_n
correction = target_ratio[h] / max(actual_ratio, EPS_DIV)
old_budget = isv[BUDGET_h_BASE + b] // for the branch this slice maps to
candidate = old_budget * correction
new_budget = old_budget + ALPHA_META * (candidate - old_budget)
new_budget = clamp(new_budget, EPS_DIV, MAX_BUDGET)
scratch[BUDGET_h_BASE + b] = new_budget
```
**Slice→branch mapping**:
| ISV slot | Branch | Slice used |
|---|---|---|
| `BUDGET_CQL_BASE+0` | direction | `dir` |
| `BUDGET_CQL_BASE+1` | magnitude | `mag` |
| `BUDGET_CQL_BASE+2` | order | `trunk` |
| `BUDGET_CQL_BASE+3` | urgency | `trunk` |
| `BUDGET_C51_BASE+0..3` | same pattern | same pattern |
Order and urgency reuse the trunk grad-norm ratio because `grad_decomp_kernel` only emits trunk/dir/mag slices. This is acceptable because:
- Order (3 actions) and urgency (3 actions) get the same gradient pressure as trunk receives;
- The TRUE per-branch grad-norm extension would require modifying `grad_decomp_kernel` to track 4 branches, which is out of scope and not load-bearing for the immediate problem (CQL dominance affects magnitude head, which we DO have a dedicated slice for).
## Cold start (Pearl A — sentinel bootstrap)
ISV slots start at sentinel 0 (per `pearl_first_observation_bootstrap`). The
key insight: `cql_norm[s]` reflects the *floor-applied* budget (the consumer
`compute_adaptive_budgets` already floors a 0-valued ISV read to
`COLD_START_FLOOR=0.02` before passing to `apply_cql_saxpy`). So the
"current effective budget" on a sentinel step is the floor, not zero.
First controller invocation seeds:
```
if cql_norm[s] < EPS_DIV || iqn_norm[s] < EPS_DIV:
// Either side is uninstrumented or warmup hasn't started.
seed = COLD_START_FLOOR // 0.02 — Invariant 1 anchor
else if old_budget < EPS_DIV: // sentinel 0 → first observation
// The norm we observed was produced under cql_budget=COLD_START_FLOOR
// (consumer-side floor). Multiplicative correction with that floor as
// the basis seeds at the value that achieves target_ratio next step.
actual_ratio = cql_norm[s] / iqn_norm[s]
correction = target_ratio[h] / max(actual_ratio, EPS_DIV)
seed = COLD_START_FLOOR * correction
seed = clamp(seed, EPS_DIV, MAX_BUDGET)
else:
// Standard EMA update path (above).
```
Same logic for C51.
**Contract**: the cold-start basis (`COLD_START_FLOOR=0.02`) must equal the
floor used by `compute_adaptive_budgets` in `fused_training.rs:3386`. If
that floor changes, this constant moves with it. Same anchor, same value,
defined once in `loss_balance_controller_kernel.cu` and referenced from a
header comment in `fused_training.rs`.
## Fold-boundary reset (Pearl A)
The slots `BUDGET_CQL_BASE+0..4` and `BUDGET_C51_BASE+0..4` are already registered in `StateResetRegistry` (`state_reset_registry.rs:585-595`) with sentinel 0. **No new registration needed.** On fold boundary the controller automatically re-runs cold-start logic.
## Smoothing (Pearl D)
The kernel writes raw new-budget values to a scratch buffer. The existing `apply_pearls_ad_kernel` then runs:
- Pearl A: detects sentinel 0 in source ISV, replaces directly (skips smoothing on first observation).
- Pearl D: Wiener-optimal EMA based on `(diff_var / (diff_var + sample_var))`.
This is the standard SP4 pattern (`pearl_first_observation_bootstrap` + `pearl_wiener_optimal_adaptive_alpha`). The new controller integrates with no new smoothing kernel.
## Pearl 2 changes
`pearl_2_budget_kernel.cu` is modified to **stop writing** to `scratch[BUDGET_CQL_BASE..]`, `scratch[BUDGET_C51_BASE..]`, `scratch[BUDGET_ENS_BASE..]`. It still writes:
- `BUDGET_IQN_BASE+0..4` — IQN budget (reference, needs to stay stable)
- `FLATNESS_BASE+0..4` — consumed by NoisyNet σ pearl
Pearl 2 launch order in `training_loop.rs:3704-3710` is unchanged. The new controller launches **after** all `grad_decomp_launch_*` calls (so iqn/cql/c51 norms are populated) and **before** `apply_pearls_ad_kernel` (so smoothing applies to the new outputs).
The `compute_adaptive_budgets` function in `fused_training.rs:3376` is unchanged — it reads ISV slots regardless of who wrote them.
## Stale doc cleanup (per `feedback_trust_code_not_docs`)
The docstring at `fused_training.rs:3409` claims:
```rust
/// B4/G5: Last adaptive CQL gradient budget (0.10×(1−regime)×health).
```
This formula is not implemented anywhere — the actual code at `fused_training.rs:3386` just floors to 0.02. The "regime_stability allocator" referenced in `state_reset_registry.rs:588` does not exist. **Both docstrings get rewritten** to describe the new controller (or deleted if redundant).
## Files touched
| File | Change |
|---|---|
| **NEW** `crates/ml/src/cuda_pipeline/loss_balance_controller_kernel.cu` | ~140 LOC kernel (1×8 grid, 2 heads × 4 branches) |
| `crates/ml/src/cuda_pipeline/pearl_2_budget_kernel.cu` | drop CQL/C51/ENS writes; keep IQN + flatness |
| `crates/ml/src/trainers/dqn/fused_training.rs` | add launch site after grad_decomp; update docstring at line 3409 |
| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | kernel slot in trainer struct + load + capture path; new `launch_loss_balance_controller` fn |
| `crates/ml/build.rs` | cubin manifest entry for new kernel |
| `docs/dqn-wire-up-audit.md` | Fix 31 entry — what landed, what was deleted, what's new |
| `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/pearl_loss_balance_controller.md` | new memory pearl |
| `~/.claude/projects/-home-jgrusewski-Work-foxhunt/memory/MEMORY.md` | index entry for the new pearl |
## Invariant-1 anchors (per `feedback_isv_for_adaptive_bounds`)
These are constants because they encode structural design choices, not tuned values that should adapt:
| Constant | Value | Justification |
|---|---|---|
| `TARGET_CQL_OVER_IQN` | 2.0 | CQL is regularizer; carries 67% of (CQL+IQN) signal weight by design |
| `TARGET_C51_OVER_IQN` | 1.0 | C51 and IQN are both distributional; equal weight |
| `ALPHA_META` | 0.01 | Slow EMA; matches `pearl_adaptive_moe_lambda` baseline |
| `EPS_DIV` | 1e-8 | Numerical-stability anchor; same as Pearl 2 |
| `COLD_START_FLOOR` | 0.02 | Cold-start budget when grad signal unavailable; matches existing fused_training floor |
| `MAX_BUDGET` | 1.0 | Upper bound; budget is a multiplier in [0, 1] |
## Acceptance criteria
1. **Code lands clean**: `SQLX_OFFLINE=true cargo check -p ml --offline` and `SQLX_OFFLINE=true cargo build -p ml --release --offline --features cuda` both green at branch HEAD.
2. **Local RTX 3050 Ti smoke**: `multi_fold_convergence` passes with non-trivial `cql_budget_per_branch` drift away from 0.02.
3. **L40S 5-epoch smoke**: `grad_split_bwd cql/iqn` ratio at HEALTH_DIAG[4] is within `[1.0, 4.0]` (target 2.0; allow 2× tolerance during transient). Same for `c51/iqn` within `[0.5, 2.0]` (target 1.0).
4. **L40S 50-epoch full validation**: magnitude Q-spread `q_full − q_quarter > 0.01` by ep20 (currently <0.003 at ep7). `eval_dist` shows non-trivial Half/Full mass (currently `eq=1.0 eh=0 ef=0`).
5. **No regressions**: smoke `magnitude_distribution` still passes. `label_scale` stays in [0.03, 0.6] range (Bug-1 fix preserved).
## Risk + rollback
**Risk 1: feedback loop instability**. If `α=0.01` is too aggressive, budget oscillation could cascade into Q-value oscillation. Mitigation: clamp range `[EPS_DIV, MAX_BUDGET=1.0]` prevents catastrophic drift. Pearl D smoothing further dampens.
**Risk 2: target ratio is wrong**. If CQL=2× IQN turns out to suppress IQN learning, magnitude differentiation may fail in different ways. Mitigation: target ratios are Invariant-1 anchors but the framework supports trivially changing them; smoke metrics will surface the issue at ep5.
**Risk 3: cold-start interaction with fold boundary**. The `cql_norm < EPS_DIV` check seeds with `COLD_START_FLOOR=0.02` — same as current behavior. If fold boundary lands a step where the gradient hasn't been computed yet (chunk boundaries), seed is benign. Mitigation: `EPS_DIV` is `1e-8`, well below any real gradient signal.
**Rollback**: revert the commit. Pearl 2's CQL/C51/ENS writes restored automatically (preserved as removed-then-re-added in single revert). No data migration concerns.
## Out-of-scope but related
- **C51 budget below floor mystery**: the `c51_budget_per_branch [dir=0.0500 …]` flat-at-floor observation suggests Pearl 2's flatness logic produces values below the floor. This is independent of SP7 and resolved automatically by SP7 since Pearl 2 stops writing to `BUDGET_C51_BASE`. (Pearl 2's flatness signal is preserved at `FLATNESS_BASE` for NoisyNet σ.)
- **Per-branch grad_decomp**: extending `grad_component_delta_norm` from 3-slice (trunk/dir/mag) to 4-branch (trunk/dir/mag/ord/urg) would let SP7 use dedicated grad-norms for ord and urg branches instead of trunk fallback. Worth doing later as a generalization, but the current trunk fallback is the right MVP.
- **Wiener-optimal α for the controller itself**: per `pearl_wiener_optimal_adaptive_alpha`, α could be derived from `diff_var / (diff_var + sample_var)`. Skipping for MVP — fixed `α=0.01` is consistent with `pearl_adaptive_moe_lambda` baseline. Phase 2 if the smoke shows oscillation.
## Memory pearls created/updated
**NEW** `pearl_loss_balance_controller.md`:
> When two loss heads contribute to the same parameter buffer at structurally
> different magnitudes (e.g., CQL's softmax-of-logsumexp emits gradient on
> all actions × atoms while IQN only emits on the action taken), a flatness-
> driven budget controller cannot balance them — it operates on Q-value
> distribution shape, not on observed gradient magnitudes. The fix is an
> outcome-driven controller that targets a fixed gradient-norm RATIO between
> the heads and adapts the multiplier to achieve it. Reference head (IQN)
> stays under existing flatness control; managed heads (CQL, C51) get
> ratio-driven budgets. Generalizes `pearl_adaptive_moe_lambda` to
> multi-head loss balancing.
**Why**: future-me dispatching a similar adaptive-coefficient task should reach for ratio-driven (not signal-driven) controllers when the loss heads have known magnitude asymmetry.
**How to apply**: when a new loss head is added to the training stack, decide whether it should be a reference (own its budget) or managed (driven by ratio to a reference). Add the controller's launch site after the head's `grad_decomp_launch_*` so the norm is populated.