fix: isolate shrink_perturb and scale_adam_momentum from graphed CUmodule
shrink_and_perturb() is called at epoch boundaries after child graphs are captured. shrink_perturb_kernel lives in the same CUmodule as scale_f32, saxpy_f32, spectral_norm, adam_update, and other graphed CUfunctions. On Hopper (sm_90), launching any CUfunction from a graphed CUmodule ungraphed corrupts the child graph kernel state → 3100ms replay. scale_adam_momentum() is called at cosine LR warm restarts (also after graph capture). scale_f32_kernel is captured in forward_child — same CUmodule violation. Fix: add shrink_perturb_ungraphed and scale_f32_ungraphed loaded from the existing ungraphed_module (separate CUmodule instance of the same DQN_UTILITY_CUBIN). Both ungraphed callers now use isolated handles. CUmodule count: 5 → 5 (ungraphed_module already existed, reused). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -642,6 +642,16 @@ pub struct GpuDqnTrainer {
|
||||
zero_kernel: CudaFunction,
|
||||
regime_scale_kernel: CudaFunction,
|
||||
shrink_perturb_kernel: CudaFunction,
|
||||
/// Ungraphed variant of shrink_perturb — from a separate CUmodule.
|
||||
/// shrink_and_perturb() runs at epoch boundaries (after graph capture).
|
||||
/// On Hopper, launching ANY CUfunction from a CUmodule that also has
|
||||
/// graph-captured CUfunctions corrupts graph kernel state (3100ms replay).
|
||||
shrink_perturb_ungraphed: CudaFunction,
|
||||
/// Ungraphed variant of scale_f32 — from a separate CUmodule.
|
||||
/// scale_adam_momentum() runs at cosine LR warm restarts (after graph capture).
|
||||
/// scale_f32_kernel is captured in forward_child; launching from the same
|
||||
/// CUmodule ungraphed corrupts the child graph on Hopper.
|
||||
scale_f32_ungraphed: CudaFunction,
|
||||
relu_mask_kernel: CudaFunction,
|
||||
spectral_norm_batched_kernel: CudaFunction,
|
||||
/// Descriptor buffer for batched spectral norm: 12 matrices x 6 u64 = 72 elements.
|
||||
@@ -1492,8 +1502,11 @@ impl GpuDqnTrainer {
|
||||
let blocks = ((self.total_params + 255) / 256) as u32;
|
||||
let cfg = LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), shared_mem_bytes: 0 };
|
||||
let m_ptr = self.m_buf.raw_ptr();
|
||||
// Use scale_f32_ungraphed — called at cosine LR warm restarts (after graph capture).
|
||||
// scale_f32_kernel is captured in forward_child; launching from the same CUmodule
|
||||
// ungraphed corrupts the child graph's kernel state on Hopper (3100ms replay).
|
||||
unsafe {
|
||||
self.stream.launch_builder(&self.scale_f32_kernel)
|
||||
self.stream.launch_builder(&self.scale_f32_ungraphed)
|
||||
.arg(&m_ptr)
|
||||
.arg(&scale)
|
||||
.arg(&n)
|
||||
@@ -4055,11 +4068,15 @@ impl GpuDqnTrainer {
|
||||
// GPU-native: single kernel generates noise + blends in one pass on params_buf.
|
||||
// params_buf[i] = alpha * params_buf[i] + (1-alpha) * N(0, sigma)
|
||||
// Branch heads [skip_start, skip_end) are excluded from perturbation.
|
||||
// Use shrink_perturb_ungraphed — this runs AFTER graph capture (epoch boundary),
|
||||
// and shrink_perturb_kernel is in the same CUmodule as graphed kernels.
|
||||
// On Hopper, launching any CUfunction from a graphed CUmodule ungraphed
|
||||
// corrupts the child graph's kernel state (3100ms replay).
|
||||
let _evt_guard = EventTrackingGuard::new(self.stream.context());
|
||||
let params_ptr = self.params_buf.raw_ptr();
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.shrink_perturb_kernel)
|
||||
.launch_builder(&self.shrink_perturb_ungraphed)
|
||||
.arg(¶ms_ptr)
|
||||
.arg(&alpha)
|
||||
.arg(&sigma)
|
||||
@@ -4106,7 +4123,7 @@ impl GpuDqnTrainer {
|
||||
// per array. Stack is set once in DQNTrainer::new() (64KB for all kernels).
|
||||
|
||||
// ── Compile utility kernels (grad_norm, adam_update, etc.) ─
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux) =
|
||||
let (grad_norm_kernel, grad_norm_finalize_kernel, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update_kernel, adam_update_post_aux, scale_f32_post_aux, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clip_grad_kernel, pad_states_kernel, saxpy_f32_kernel, scale_f32_kernel, stochastic_depth_kernel, stochastic_depth_rng_kernel, bn_tanh_backward_kernel, bn_bias_grad_kernel, bn_tanh_concat_kernel_fn, vaccine_dot_kernel, vaccine_project_kernel, vaccine_dot_finalize, causal_intervene_kernel_fn, causal_reduce_kernel_fn, causal_mean_reduce_kernel_fn, pruning_mask_kernel, pruning_compute_kernel, her_inplace_kernel, nan_check_f32_kernel, nan_check_f32_kernel_b, popart_normalize_kernel, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed) =
|
||||
compile_training_kernels(&stream, &config)?;
|
||||
|
||||
// Separate grad_norm instance for non-graph launches (outside CUDA graph).
|
||||
@@ -5720,6 +5737,8 @@ impl GpuDqnTrainer {
|
||||
zero_kernel,
|
||||
regime_scale_kernel,
|
||||
shrink_perturb_kernel: shrink_perturb,
|
||||
shrink_perturb_ungraphed,
|
||||
scale_f32_ungraphed,
|
||||
relu_mask_kernel: relu_mask_standalone,
|
||||
spectral_norm_batched_kernel: spectral_norm_kernel,
|
||||
spectral_norm_descriptors,
|
||||
@@ -10302,7 +10321,7 @@ impl GpuDqnTrainer {
|
||||
fn compile_training_kernels(
|
||||
stream: &Arc<CudaStream>,
|
||||
config: &GpuDqnTrainConfig,
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
info!(
|
||||
state_dim = config.state_dim,
|
||||
total_params = compute_total_params(config),
|
||||
@@ -10372,15 +10391,26 @@ fn compile_training_kernels(
|
||||
.map_err(|e| MLError::ModelError(format!("spectral_norm_batched load: {e}")))?;
|
||||
let clip_grad = module.load_function("dqn_clip_grad_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("dqn_clip_grad_kernel load: {e}")))?;
|
||||
// pad_states + stochastic_depth_rng run UNGRAPHED every step. They MUST be from
|
||||
// a separate CUmodule — on Hopper, launching ANY CUfunction from a CUmodule that
|
||||
// also has graph-captured CUfunctions corrupts the graph's kernel state (3100ms replay).
|
||||
// Ungraphed kernels MUST be from a separate CUmodule — on Hopper, launching ANY
|
||||
// CUfunction from a CUmodule that also has graph-captured CUfunctions corrupts
|
||||
// the graph's kernel state (3100ms replay). All kernels launched outside graph
|
||||
// capture (per-step or at epoch boundaries) must use handles from this module.
|
||||
let ungraphed_module = context.load_cubin(DQN_UTILITY_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("ungraphed utility cubin load: {e}")))?;
|
||||
let pad_states = ungraphed_module.load_function("pad_states_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("pad_states_kernel load: {e}")))?;
|
||||
let stochastic_depth_rng = ungraphed_module.load_function("stochastic_depth_rng")
|
||||
.map_err(|e| MLError::ModelError(format!("stochastic_depth_rng load: {e}")))?;
|
||||
// shrink_and_perturb() is called at epoch boundaries (after graph capture).
|
||||
// dqn_shrink_perturb_kernel lives in the same CUmodule as other graphed kernels
|
||||
// (scale_f32, saxpy_f32, spectral_norm, adam_update, etc.) — so it MUST be
|
||||
// loaded from ungraphed_module to prevent Hopper CUmodule corruption.
|
||||
let shrink_perturb_ungraphed = ungraphed_module.load_function("dqn_shrink_perturb_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("shrink_perturb_ungraphed load: {e}")))?;
|
||||
// scale_adam_momentum() is called at cosine LR warm restarts (after graph capture).
|
||||
// dqn_scale_f32_kernel is captured in forward_child — same CUmodule as above.
|
||||
let scale_f32_ungraphed = ungraphed_module.load_function("dqn_scale_f32_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("scale_f32_ungraphed load: {e}")))?;
|
||||
// stochastic_depth_scale is GRAPHED (inside forward_child via apply_regime_dropout)
|
||||
// — stays in the main module.
|
||||
let stochastic_depth = module.load_function("stochastic_depth_scale")
|
||||
@@ -10418,8 +10448,8 @@ fn compile_training_kernels(
|
||||
let popart_normalize = module.load_function("popart_normalize_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("popart_normalize_kernel load: {e}")))?;
|
||||
|
||||
info!("GpuDqnTrainer: 33 utility kernels loaded from precompiled cubin (5 CUmodules)");
|
||||
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux))
|
||||
info!("GpuDqnTrainer: 35 utility kernels loaded from precompiled cubin (5 CUmodules)");
|
||||
Ok((grad_norm, grad_norm_finalize, grad_norm_finalize_adam, grad_norm_finalize_post_aux, adam_update, adam_update_post_aux, scale_f32_post_aux, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clip_grad, pad_states, saxpy_f32, scale_f32, stochastic_depth, stochastic_depth_rng, bn_tanh_bw, bn_bias_grad, bn_tanh_concat, vaccine_dot, vaccine_project, vaccine_dot_finalize, causal_intervene, causal_reduce, causal_mean_reduce, pruning_mask_fn, pruning_compute_fn, her_inplace, nan_check_f32, nan_check_f32_b, popart_normalize, saxpy_f32_adam_grad, saxpy_f32_aux, grad_norm_standalone_post_aux, shrink_perturb_ungraphed, scale_f32_ungraphed))
|
||||
}
|
||||
|
||||
/// Load the standalone Polyak EMA kernel from precompiled cubin.
|
||||
|
||||
Reference in New Issue
Block a user