[Stacked PR 3/5] Add decoupled GDN Pallas backward pass kernel and parity tests - #5153
[Stacked PR 3/5] Add decoupled GDN Pallas backward pass kernel and parity tests#5153Rohan-Bierneni wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a local Gated Delta Net (GDN) backward kernel package using Pallas emit_pipeline with triangular inverse matrix caching, alongside a decoupled Conv1D + SiLU forward and backward implementation in JAX, and comprehensive unit tests. The review feedback identifies two key issues: an epsilon mismatch (1e-12 vs. 1e-6) when use_qk_norm_in_gdn is enabled, which could cause numerical instability, and an unconditional gradient reduction for a_log and dt_bias that fails to preserve batch dimensions for 2D or 3D inputs.
| if use_qk_norm_in_gdn: | ||
| d_q_scaled = d_q_proj * scale | ||
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) | ||
| q_unit = q_orig / r_q | ||
| d_q = (d_q_scaled - q_unit * jnp.sum(d_q_scaled * q_unit, axis=-1, keepdims=True)) / r_q | ||
|
|
||
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) | ||
| k_unit = k_orig / r_k | ||
| d_k = (d_k_proj - k_unit * jnp.sum(d_k_proj * k_unit, axis=-1, keepdims=True)) / r_k | ||
| else: | ||
| d_q = d_q_proj * scale | ||
| d_k = d_k_proj |
There was a problem hiding this comment.
There is an epsilon mismatch between the forward and backward passes when use_qk_norm_in_gdn is enabled. In the forward pass, normalizations.l2norm is called with eps=1e-6. However, in the backward pass, r_q and r_k are computed using 1e-12 as the epsilon.
While this mismatch is negligible in unit tests with random normal inputs (where the sum of squares is large), it can lead to mathematically incorrect gradients, gradient explosion, or NaNs in real training scenarios—especially for padded/masked tokens or during early initialization where values can be extremely small or zero. Changing 1e-12 to 1e-6 ensures mathematical consistency and training stability.
| if use_qk_norm_in_gdn: | |
| d_q_scaled = d_q_proj * scale | |
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) | |
| q_unit = q_orig / r_q | |
| d_q = (d_q_scaled - q_unit * jnp.sum(d_q_scaled * q_unit, axis=-1, keepdims=True)) / r_q | |
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) | |
| k_unit = k_orig / r_k | |
| d_k = (d_k_proj - k_unit * jnp.sum(d_k_proj * k_unit, axis=-1, keepdims=True)) / r_k | |
| else: | |
| d_q = d_q_proj * scale | |
| d_k = d_k_proj | |
| if use_qk_norm_in_gdn: | |
| d_q_scaled = d_q_proj * scale | |
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-6) | |
| q_unit = q_orig / r_q | |
| d_q = (d_q_scaled - q_unit * jnp.sum(d_q_scaled * q_unit, axis=-1, keepdims=True)) / r_q | |
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-6) | |
| k_unit = k_orig / r_k | |
| d_k = (d_k_proj - k_unit * jnp.sum(d_k_proj * k_unit, axis=-1, keepdims=True)) / r_k | |
| else: | |
| d_q = d_q_proj * scale | |
| d_k = d_k_proj |
| d_a_log_reduced = jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(a_log.dtype) | ||
| d_dt_bias_reduced = jnp.sum(d_dt_bias_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(dt_bias.dtype) |
There was a problem hiding this comment.
The forward path of _pallas_gdn_bwd_kernel_single_group explicitly supports 1D, 2D, and 3D shapes for a_log and dt_bias. However, the backward path unconditionally reduces the gradients over axis=(0, 1) (batch and chunk dimensions):
d_a_log_reduced = jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=(0, 1))If a_log or dt_bias is 2D (e.g., batch-dependent), this unconditional reduction over axis=0 will destroy the batch dimension, leading to shape mismatches or incorrect gradient updates in JAX. We should dynamically reduce the gradients based on the input dimensions of a_log and dt_bias.
| d_a_log_reduced = jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(a_log.dtype) | |
| d_dt_bias_reduced = jnp.sum(d_dt_bias_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(dt_bias.dtype) | |
| if a_log.ndim == 1: | |
| d_a_log_reduced = jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(a_log.dtype) | |
| elif a_log.ndim == 2: | |
| d_a_log_reduced = jnp.sum(d_a_log_chunks[..., 0, :num_v_heads], axis=1).astype(a_log.dtype) | |
| else: | |
| d_a_log_reduced = d_a_log_chunks[..., 0, :num_v_heads].astype(a_log.dtype) | |
| if dt_bias.ndim == 1: | |
| d_dt_bias_reduced = jnp.sum(d_dt_bias_chunks[..., 0, :num_v_heads], axis=(0, 1)).astype(dt_bias.dtype) | |
| elif dt_bias.ndim == 2: | |
| d_dt_bias_reduced = jnp.sum(d_dt_bias_chunks[..., 0, :num_v_heads], axis=1).astype(dt_bias.dtype) | |
| else: | |
| d_dt_bias_reduced = d_dt_bias_chunks[..., 0, :num_v_heads].astype(dt_bias.dtype) |
Rohan-Bierneni
changed the title
9238216 to
d4ca233
Compare
…rity tests - Implement decoupled GDN backward pass using Pallas emit_pipeline (pallas_gdn_bwd_kernel) and custom VJP (gdn_decoupled_conv1d). - Execute Conv1D backward immediately following Pallas via unrolled native JAX shift-multiply, eliminating live range interference and vector register spills. - Provide comprehensive analytical gradient parity test suite (gdn_bwd_pallas_test.py) validating exact gradient match against pure JAX autodiff across varying sequence lengths, batch sizes, and head configurations.
d4ca233 to
c47f1a3
Compare
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
Stacked PR Chain
rbierneni-gdn-1-ci-hygienemainrbierneni-gdn-2-fwd-kernelrbierneni-gdn-1-ci-hygienerbierneni-gdn-3-bwd-kernelrbierneni-gdn-2-fwd-kernelrbierneni-gdn-4-model-integrationrbierneni-gdn-3-bwd-kernelrbierneni-gdnv3-bwdrbierneni-gdn-4-model-integrationDescription
This is PR 3 of 5 in the stacked series enabling the Pallas Gated Delta Net (GDN) backward pass kernel in MaxText.
This PR introduces the decoupled GDN Pallas backward pass kernel and complete analytical gradient parity testing:
gdn_bwd_pallas.py):pallas_gdn_bwd_kernel: Executes 40+ GDN adjoint matrix recurrences viapltpu.emit_pipelinewith vectorized head processing and zero HBM intermediate spills.conv1d_silu_bwd: Executes immediately after Pallas via unrolled native JAX shift-multiply. Decoupling Conv1D backward avoids binding Conv1D gradient live ranges with GDN adjoint matrix state buffers, eliminates vector register spills, and keeps peak VMEM well within hardware boundaries.gdn_decoupled_conv1d: Custom VJP unifying forward execution (reusinggdn_decoupled_conv1d(rather thanfused) to accurately convey that Conv1D is decoupled in native JAX rather than fused with GDN recurrence in Pallas.gdn_kerneland backward-compatible aliases.tests/unit/gdn_bwd_pallas_test.py):Files Changed
src/maxtext/models/kernels/gdn/__init__.py: Package export forgdn_decoupled_conv1d,gdn_kernel,pallas_gdn_bwd_kernel.src/maxtext/models/kernels/gdn/gdn_bwd_pallas.py: Decoupled backward pass kernel and custom VJP implementation (gdn_decoupled_conv1d).tests/unit/gdn_bwd_pallas_test.py: Full analytical gradient parity test suite against JAX autodiff.Tests
pre-commit run --files ...(all hooks passed:codespell,pylint,pyink).pytest tests/unit/gdn_bwd_pallas_test.pyvalidates analytical gradient equivalence to machine precision (13/13 unit tests passed in 137s).Checklist