Test Fused GDN Backward Pass Kernel w/ MaxText - #5098
Conversation
- Port Tokamax GDN forward kernel package under src/maxtext/models/kernels/gdn/ - Add hybrid_bwd_analytical_pipeline.py with analytical backward pass - Expose use_gdn_kernel in base.yml and types.py (default: false) - Integrate use_gdn_kernel in Qwen3NextGatedDeltaNet (qwen3.py and qwen3_5.py) - Port unit test hybrid_bwd_analytical_pipeline_test.py
… test - Fix sublane tiling calculation in pallas_mosaic_tpu_v2_gmm_kernel.py: ensure minimum size 16 - Enforce symmetric common sublane in pallas_mosaic_tpu_v2_tgmm_kernel.py for Qwen3.5 compatibility - Add test_qwen3_5_gdn_kernel in tests/unit/train_compile_test.py for AOT compilation verification on tpu7x-512
![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 fused analytical Pallas Gated Delta Net (GDN) kernel and backward pipeline, integrating it into the Qwen3 model. Key changes include the addition of the hybrid_bwd_analytical_pipeline and a suite of supporting local GDN kernel files to handle causal Conv1D and Gated Delta Rule operations with cached triangular inverse matrices. Additionally, sublane tiling logic is adjusted in the GMM and TGMM kernels. Feedback on the changes highlights a numerical inconsistency where an epsilon of 1e-12 is used for L2 normalization in the backward pass instead of the 1e-6 used in the forward pass, as well as an opportunity to simplify the sublane alignment assignment in the TGMM kernel.
|
|
||
| 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) |
There was a problem hiding this comment.
The epsilon value 1e-12 used here for calculating the norm of q is inconsistent with the value 1e-6 used in the forward pass (l2norm at line 521). For numerical stability and correctness of the gradient, the same epsilon value should be used in both forward and backward computations.
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-12) | |
| r_q = jnp.sqrt(jnp.sum(q_orig**2, axis=-1, keepdims=True) + 1e-6) |
| - 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) |
There was a problem hiding this comment.
Similar to the calculation for r_q, the epsilon value 1e-12 used here for r_k is inconsistent with the 1e-6 used in the forward pass. This should be corrected to 1e-6 to ensure numerical consistency.
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-12) | |
| r_k = jnp.sqrt(jnp.sum(k_orig**2, axis=-1, keepdims=True) + 1e-6) |
| common_sublane = min(size_lhs_sublane, size_rhs_sublane) | ||
| size_lhs_sublane = common_sublane | ||
| size_rhs_sublane = common_sublane |
There was a problem hiding this comment.
The logic to ensure size_lhs_sublane and size_rhs_sublane are equal can be simplified into a single line, improving readability.
| common_sublane = min(size_lhs_sublane, size_rhs_sublane) | |
| size_lhs_sublane = common_sublane | |
| size_rhs_sublane = common_sublane | |
| size_lhs_sublane = size_rhs_sublane = min(size_lhs_sublane, size_rhs_sublane) |
- Add src/maxtext/configs/models/qwen3.5-tiny.yml for fast 132M parameter testing - Register qwen3.5-tiny in ModelName Literal and valid_mm_models in types.py - Register qwen3.5-tiny in HF_IDS in globals.py - Register qwen3.5-tiny in multimodal embedding whitelist in decoders.py and nnx_decoders.py
1 participant
Description
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.