diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 90a3f374d..b3fb66f8b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2131,6 +2131,39 @@ pub struct GpuDqnTrainer { bw_d_glu_value: CudaSlice, /// GLU backward scratch: d_gate_pre [B, AH] — f32 (reused across branches) bw_d_glu_gate: CudaSlice, + + // ── GRN trunk backward scratch (Plan 4 Task 2c.3c.3) ───────────── + // 4 buffers per GRN block × 2 blocks (h_s1 + h_s2) = 8 total. + // Consumed by `apply_grn_trunk_backward_raw` helper (Task 2c.3c.4). + // Each pair of buffers per block: + // d_pre_LN — phase1 input (d_glu_out + d_residual aliased) [B, hidden] + // d_linear_b_out — phase1 output for caller's Linear_b backward GEMM [B, 2*hidden] + // d_elu_out — caller's Linear_b backward output, phase2 input [B, hidden] + // d_linear_a_out — phase2 output for caller's Linear_a backward GEMM [B, hidden] + /// h_s2 GRN d_pre_LN scratch [B, SH2] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s2_d_pre_ln: CudaSlice, + /// h_s2 GRN d_linear_b_out scratch [B, 2*SH2] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s2_d_linear_b_out: CudaSlice, + /// h_s2 GRN d_elu_out scratch [B, SH2] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s2_d_elu_out: CudaSlice, + /// h_s2 GRN d_linear_a_out scratch [B, SH2] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s2_d_linear_a_out: CudaSlice, + /// h_s1 GRN d_pre_LN scratch [B, SH1] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s1_d_pre_ln: CudaSlice, + /// h_s1 GRN d_linear_b_out scratch [B, 2*SH1] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s1_d_linear_b_out: CudaSlice, + /// h_s1 GRN d_elu_out scratch [B, SH1] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s1_d_elu_out: CudaSlice, + /// h_s1 GRN d_linear_a_out scratch [B, SH1] — wired by Task 2c.3c.4 + #[allow(dead_code)] + bw_grn_h_s1_d_linear_a_out: CudaSlice, // ── Expected Q-value kernel (ad-hoc validation, not captured in CUDA Graph) ─ /// Converts C51 value+advantage logits → expected Q-values (validation path). /// Input: on_v_logits_buf [B, NA], on_b_logits_buf [B, (B0+B1+B2+B3)*NA] @@ -8149,6 +8182,30 @@ impl GpuDqnTrainer { alloc_backward_scratch(&stream, &config) .map_err(|e| MLError::ModelError(format!("backward scratch alloc: {e}")))?; + // ── GRN trunk backward scratch (Plan 4 Task 2c.3c.3) ──────── + // 4 buffers per GRN block × 2 blocks = 8 buffers total. Consumed by + // 2c.3c.4's `apply_grn_trunk_backward_raw` helper. Sizes mirror the + // forward trunk hidden dims (SH1 / SH2) and use `config.batch_size` + // as the batch dim, matching `alloc_backward_scratch` convention. + // d_linear_b_out is [B, 2*hidden] because GLU's pre-activation has + // 2× the hidden width (value path + gate_pre path concatenated). + let bw_grn_h_s2_d_pre_ln = stream.alloc_zeros::(b * config.shared_h2) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s2_d_pre_ln: {e}")))?; + let bw_grn_h_s2_d_linear_b_out = stream.alloc_zeros::(b * 2 * config.shared_h2) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s2_d_linear_b_out: {e}")))?; + let bw_grn_h_s2_d_elu_out = stream.alloc_zeros::(b * config.shared_h2) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s2_d_elu_out: {e}")))?; + let bw_grn_h_s2_d_linear_a_out = stream.alloc_zeros::(b * config.shared_h2) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s2_d_linear_a_out: {e}")))?; + let bw_grn_h_s1_d_pre_ln = stream.alloc_zeros::(b * config.shared_h1) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s1_d_pre_ln: {e}")))?; + let bw_grn_h_s1_d_linear_b_out = stream.alloc_zeros::(b * 2 * config.shared_h1) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s1_d_linear_b_out: {e}")))?; + let bw_grn_h_s1_d_elu_out = stream.alloc_zeros::(b * config.shared_h1) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s1_d_elu_out: {e}")))?; + let bw_grn_h_s1_d_linear_a_out = stream.alloc_zeros::(b * config.shared_h1) + .map_err(|e| MLError::ModelError(format!("alloc bw_grn_h_s1_d_linear_a_out: {e}")))?; + let batch_bytes = (b * ml_core::state_layout::STATE_DIM * 2 + b * 4 + b * config.shared_h1 @@ -9615,6 +9672,15 @@ impl GpuDqnTrainer { bw_d_h_b3, bw_d_glu_value, bw_d_glu_gate, + // GRN trunk backward scratch (Plan 4 Task 2c.3c.3) — wired by 2c.3c.4 + bw_grn_h_s2_d_pre_ln, + bw_grn_h_s2_d_linear_b_out, + bw_grn_h_s2_d_elu_out, + bw_grn_h_s2_d_linear_a_out, + bw_grn_h_s1_d_pre_ln, + bw_grn_h_s1_d_linear_b_out, + bw_grn_h_s1_d_elu_out, + bw_grn_h_s1_d_linear_a_out, expected_q_kernel, q_stats_kernel, q_stats_per_branch_kernel, diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 1186cdc6d..06b0d6972 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -289,6 +289,8 @@ Plan 4 Task 2c.3c.1 (2026-04-24): `GrnBlock::backward_raw_phase1` + `backward_ra Plan 4 Task 2c.3c.2 (2026-04-24): backward infrastructure additions — two helper methods that 2c.3c.4's GRN trunk backward wire-up requires. (1) `CublasBackwardSet::launch_dw_only_no_bias` — variant of `launch_dw_only` that runs only the cuBLAS dW GEMM and skips the 2-phase bias-grad reduction. The h_s1 GRN block's `Linear_residual` projection has no bias term (matches `gpu_grn.rs::has_residual_projection=true`'s 7-tensor layout: w_a, b_a, w_b, b_b, w_residual, gamma, beta — note no `b_residual`); calling the existing `launch_dw_only` with `db=0u64` would dereference NULL inside `bias_grad_reduce_f32_p1`. Same row-major→col-major sgemm convention as `launch_dw_only` (label `"dW_only_no_bias"` for cublasLt diagnostics), `pub(crate)` visibility matching its sibling. +35 LOC. (2) `CublasGemmSet::saxpy_inplace(stream, n, alpha, x_ptr, y_ptr)` — element-wise `y[i] += alpha * x[i]` for the h_s2 GRN's identity-residual gradient accumulation: after Linear_a_h_s2's backward writes `d_h_s1 = d_linear_a_h_s2 @ W_a_h_s2` (overwrite, β=0), the residual path needs `d_h_s1 += d_pre_ln_h_s2` to capture the identity-residual gradient. Implementation chose **option 2** (existing `dqn_saxpy_f32_kernel` from `dqn_utility_kernels.cu`) over option 1 (cuBLAS `cublasSaxpy_v2` legacy handle): the legacy cuBLAS handle is per-stream-bound state on `PerStreamCublasHandles`, so adopting it for backward saxpy would require either re-binding (race risk vs forward's binding) or a new shared handle slot — versus the existing kernel which is already loaded by `GpuExperienceCollector` for IQR/ensemble-variance Q-bonus saxpy and is graph-capturable via the standard `launch_builder` path. New `saxpy_kernel: CudaFunction` field on `CublasGemmSet`, loaded in `CublasGemmSet::new` from `DQN_UTILITY_CUBIN` (same cubin already include_bytes!'d by `gpu_dqn_trainer.rs`); +66 LOC including field, init block, and the `pub fn saxpy_inplace` method (256 threads/block, `ceil(n/256)` blocks, no grid-stride needed — kernel is already bounded by `if (i < n)`). **Additive only — ZERO production callers in this commit**; both methods sit dead-code until 2c.3c.4's wire-up commit. cargo check clean at 11 warnings (baseline preserved); cargo build unchanged — no new `.cu` files (saxpy uses the existing `dqn_saxpy_f32_kernel` symbol). No new module / kernel / ISV slot / Orphan row. +Plan 4 Task 2c.3c.3 (2026-04-24): GRN trunk backward scratch buffer allocation on `GpuDqnTrainer`. 8 new `CudaSlice` device buffers added (4 per GRN block × 2 blocks for h_s1 + h_s2): `bw_grn_h_s2_d_pre_ln` [B, SH2], `bw_grn_h_s2_d_linear_b_out` [B, 2*SH2], `bw_grn_h_s2_d_elu_out` [B, SH2], `bw_grn_h_s2_d_linear_a_out` [B, SH2], plus the four h_s1 mirrors at [B, SH1] (and 2*SH1 for the linear_b_out scratch — GLU's pre-activation has 2× hidden width because value-path + gate_pre are concatenated). The `d_pre_LN` buffer aliases both `d_glu_out` and `d_residual` (pure pass-through under `GrnBlock::backward_raw_phase1`'s pointer aliasing convention from 2c.3c.1); `d_linear_b_out` carries the phase1 → caller's Linear_b backward GEMM gradient; `d_elu_out` carries Linear_b backward's output into phase2; `d_linear_a_out` carries phase2's output to the caller's Linear_a backward GEMM. Allocated alongside the existing `alloc_backward_scratch` call in the trainer constructor (mirrors that helper's `stream.alloc_zeros::(size).map_err(...)?` pattern). Total footprint at SH1=SH2=256, B=512: 8 × 5 × 256 × 4B = 5.0 MB per fold (negligible against existing ~32 MB per-branch cuBLAS workspaces). Each field is `#[allow(dead_code)]` until 2c.3c.4's `apply_grn_trunk_backward_raw` helper consumes them — the comment `wired by Task 2c.3c.4` annotates each suppression. **Additive only — ZERO production callers in this commit**; the three backward panic gates (`backward_full`, `apply_iqn_trunk_gradient`, `apply_ensemble_diversity_backward`) remain in place until 2c.3c.4. Adam state for the new GRN tensors is auto-allocated since `m_buf`/`v_buf` are sized to `total_params + cutlass_tile_pad` which already includes the 13 GRN tensors after 2c.3a. cargo check clean at 11 warnings (baseline preserved); cargo build compiles 80 cubins (unchanged — no new kernel). +57 LOC. No new module / kernel / ISV slot / Orphan row. + Plan 4 Task 2c.3a (2026-04-24): GRN trunk param-tensor reshuffle (build-clean, runtime-broken intentionally). `compute_param_sizes()` flat layout reshuffled from 86 → 95 tensors: the 4 legacy trunk Linear tensors (`w_s1`, `b_s1`, `w_s2`, `b_s2` at indices [0..4)) are replaced with **13 GRN tensors** at indices [0..13) — 7 for h_s1 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `w_residual`, `gamma`, `beta`; SD→SH1 with `Linear_residual` GEMM since dimensions differ) plus 6 for h_s2 GRN block (`w_a`, `b_a`, `w_b`, `b_b`, `gamma`, `beta`; SH1→SH2 with identity residual since SH1==SH2). Every tensor index ≥ 4 in the OLD layout shifts +9 in the NEW layout. `NUM_WEIGHT_TENSORS` 86→95, `FIRST_ISV_TENSOR` 68→77. `layout_fingerprint_seed()` updated to mirror the new tensor names + positions; new `LAYOUT_FINGERPRINT_CURRENT = 0xcf3a24b0a1f70057` (was `0xa504d3c2f275b8af`). All 93 `padded_byte_offset` call sites + ancillary index references migrated in lockstep per `feedback_no_partial_refactor.md`. `xavier_init_params_buf` extended with init paths for the 13 new tensors: Linear matrices (`w_a`, `w_b`, `w_residual`) get Xavier; LayerNorm γ initialised to 1.0; β + biases stay zero. Trunk-forward / trunk-backward / IQN-trunk / ensemble-diversity-backward callers gated by explicit `panic!("Plan 4 Task 2c.3a: trunk param layout migrated to GRN, …")` at function entry — `BatchedForward::encoder_forward_only`, `BatchedForward::forward_target_raw`, `BatchedForward::forward_online_f32`, `BatchedBackward::backward_full`, `GpuDqnTrainer::apply_iqn_trunk_gradient`, `GpuDqnTrainer::apply_ensemble_diversity_backward`. This makes accidental runtime execution fail loudly rather than producing silent garbage from running legacy `Linear→ReLU→Linear` GEMMs against GRN-shaped param tensors. Spectral-norm descriptor (13 matrices) keeps slots [0]/[1] mapped to `w_a_h_s1`/`w_a_h_s2` (shapes match legacy W_s1/W_s2 exactly — the GRN's first Linear has the same shape as the original Linear), so the existing spectral-norm constraint transfers cleanly to the GRN's first Linear; Linear_b / Linear_residual are NOT yet spectral-normed (Task 2c.3c follow-up). **Smoke intentionally NOT run** — build-clean is the validation; runtime would hit the panic at the first encoder forward (the desired behaviour, no need to verify explicitly). Task 2c.3b swaps in the GRN forward (gpu_grn::GrnBlock::forward) at all panic-gated forward callers in place; Task 2c.3c does the same for backward callers and runs the smoke test. cargo check clean at 11 warnings (baseline preserved); cargo build compiles all 58 cubins. No new module / kernel / ISV slot in this commit — pure structural reshuffle + assertion gates. | Classification | Count |