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

Revert "fix(dqn): isolate OFI bias-grad CUfunctions in dedicated CUmodule"

Browse Source 
This reverts commit aff69261e1.
This commit is contained in:
jgrusewski
2026-04-27 21:56:15 +02:00
parent aff69261e1
commit 5857f5147a
2 changed files with 2 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -3090,14 +3090,6 @@ pub struct GpuDqnTrainer {
denoise_bias_grad_p1_kernel: CudaFunction,
/// Kernel handle: denoise_bias_grad_p2.
denoise_bias_grad_p2_kernel: CudaFunction,
/// Dedicated CUfunction handles for OFI bias-grad reduction.
/// Loaded from a SEPARATE CUmodule from the denoise variant — each
/// child graph (post_aux, adam_grad) must hold its own CUfunction
/// reference, otherwise the driver's per-node handle validation on
/// Ada/Hopper trips CUDA_ERROR_ILLEGAL_ADDRESS at kernel launch
/// (same pattern as graph_safe_copy_kernel below).
ofi_embed_bias_grad_p1_kernel: CudaFunction,
ofi_embed_bias_grad_p2_kernel: CudaFunction,
// ── xLSTM temporal context (mLSTM cell) ────────────────────────────
/// [528] 6 weight matrices × [11, 8]: W_k, W_v, W_q, W_i, W_f, W_o
@@ -5188,7 +5180,7 @@ impl GpuDqnTrainer {
let grad_b_ptr = self.ofi_embed_grad_b.raw_ptr();
let dim_i32 = OFI_EMBED_DIM as i32;
unsafe {
self.stream.launch_builder(&self.ofi_embed_bias_grad_p1_kernel)
self.stream.launch_builder(&self.denoise_bias_grad_p1_kernel)
.arg(&d_combined_ptr)
.arg(&partials_ptr)
.arg(&dim_i32)
@@ -5201,7 +5193,7 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p1: {e}")))?;
}
unsafe {
self.stream.launch_builder(&self.ofi_embed_bias_grad_p2_kernel)
self.stream.launch_builder(&self.denoise_bias_grad_p2_kernel)
.arg(&partials_ptr)
.arg(&grad_b_ptr)
.arg(&dim_i32)
@@ -11105,16 +11097,6 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p1 load: {e}")))?;
let denoise_bias_grad_p2_kernel = cpbi_module_denoise_cublas.load_function("denoise_bias_grad_p2")
.map_err(|e| MLError::ModelError(format!("denoise_bias_grad_p2 load: {e}")))?;
// Dedicated CUmodule for OFI bias-grad reduction (Ada/Hopper child-graph
// CUfunction isolation — same pattern as graph_safe_copy below).
let ofi_embed_bias_grad_module = stream.context().load_cubin(EXPECTED_Q_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad module load: {e}")))?;
let ofi_embed_bias_grad_p1_kernel = ofi_embed_bias_grad_module
.load_function("denoise_bias_grad_p1")
.map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p1 load: {e}")))?;
let ofi_embed_bias_grad_p2_kernel = ofi_embed_bias_grad_module
.load_function("denoise_bias_grad_p2")
.map_err(|e| MLError::ModelError(format!("ofi_embed_bias_grad_p2 load: {e}")))?;
// cuBLAS GEMM descriptors for denoise MLP backward
let denoise_lt_handle = shared_cublas.lt_handle.0;
@@ -12055,8 +12037,6 @@ impl GpuDqnTrainer {
denoise_skip_conn_kernel,
denoise_bias_grad_p1_kernel,
denoise_bias_grad_p2_kernel,
ofi_embed_bias_grad_p1_kernel,
ofi_embed_bias_grad_p2_kernel,
qlstm_weights,
qlstm_c,
qlstm_n,

12
docs/dqn-wire-up-audit.md
View File

@@ -858,15 +858,3 @@ architecture pretended to do. Workspace clean: 0 errors across all crates, tests
Smoke validates no regression: 3/3 folds passed (405s), gate utilization
preserved (util=0.119,0.119,0.169,...), HEALTH_DIAG aux_moe line emits,
all fold checkpoints written.
OFI bias-grad CUfunction isolation (2026-04-27): fixed
CUDA_ERROR_ILLEGAL_ADDRESS crash in launch_ofi_embed_backward on L40S
(workflow train-multi-seed-tdjtr, fold 0). Root cause: shared
CUfunction handles for denoise_bias_grad_p1/p2 across two captured
child graphs (post_aux, adam_grad) trip Ada/Hopper driver per-node
handle validation when the partials buffers land far apart in VRAM.
Pre-existing latent bug, surfaced by MoE allocator-state shift.
Fix: dedicated CUmodule for OFI bias-grad kernels, mirroring the
graph_safe_copy_kernel precedent (line ~11523). Kernel bytecode
identical; constructor adds ~2ms; zero runtime cost. Smoke validates
no regression at B=64.