feat(dqn-v2): Plan 4 Task 2c.3c.3 — allocate 8 GRN backward scratch buffers
Plan 4 Task 2c.3c.3. Additive only — buffers will be consumed by
2c.3c.4's apply_grn_trunk_backward_raw helper.
8 device scratch buffers (4 per GRN block × 2 blocks):
- grn_h_s{1,2}_d_pre_ln — [B, hidden] aliases d_glu_out, d_residual
- grn_h_s{1,2}_d_linear_b_out — [B, 2*hidden]
- grn_h_s{1,2}_d_elu_out — [B, hidden]
- grn_h_s{1,2}_d_linear_a_out — [B, hidden]
Total ~B × 5 × hidden f32 per block × 2 blocks ≈ B × 5120 floats
(20 KB per sample at SH1=SH2=256). Negligible against existing
~32 MB per-branch cuBLAS workspaces.
Allocated in the trainer's alloc_backward_scratch path alongside
existing per-step backward scratch. dead_code allow on each field
until 2c.3c.4 wires them in.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -2131,6 +2131,39 @@ pub struct GpuDqnTrainer {
|
||||
bw_d_glu_value: CudaSlice<f32>,
|
||||
/// GLU backward scratch: d_gate_pre [B, AH] — f32 (reused across branches)
|
||||
bw_d_glu_gate: CudaSlice<f32>,
|
||||
|
||||
// ── 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
/// 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<f32>,
|
||||
// ── 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::<f32>(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::<f32>(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::<f32>(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::<f32>(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::<f32>(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::<f32>(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::<f32>(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::<f32>(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,
|
||||
|
||||
@@ -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<f32>` 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::<f32>(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 |
|
||||
|
||||
Reference in New Issue
Block a user