[Stacked PR 4/5] Enable GDN backward pass kernel and hybrid precision in Qwen3 model - #5154
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 configuration options and support for the GDN Pallas kernel (use_gdn_kernel) along with customizable precision types (gdn_state_dtype and gdn_decay_dtype) in the Qwen3 model. Feedback on these changes highlights two key areas for improvement: first, under ShardMode.EXPLICIT, manual resharding of inputs is required for the Pallas GDN kernel to prevent compilation or runtime errors due to layout mismatches in shard_map; second, the dtype resolution logic for GDN decay and state parameters can be simplified and made more robust by utilizing jnp.dtype directly instead of manual string checks.
| qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) | ||
| b_a_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) | ||
| conv_state_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules) | ||
| recurrent_state_pspec = logical_to_mesh_axes((KV_BATCH, None, None, None), mesh=self.mesh, rules=logical_rules) |
There was a problem hiding this comment.
Under ShardMode.EXPLICIT, shard_map manualizes the mesh axes and does not automatically insert resharding/all-gathers for operands whose layouts differ from in_specs. Since the Pallas GDN kernel expects the sequence and head dimensions to be local (replicated) on each device, calling gdn_step_fn directly without resharding will cause compilation/runtime errors.
We should adjust the PartitionSpecs using remove_incompatible_mesh_axes_from_partition_spec and manually reshard the inputs (qkv, b, a, conv_state_arg, and recurrent_state_arg) to match their respective PartitionSpecs when cfg.shard_mode == ShardMode.EXPLICIT.
qkv_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules)
b_a_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules)
conv_state_pspec = logical_to_mesh_axes((KV_BATCH, None, None), mesh=self.mesh, rules=logical_rules)
recurrent_state_pspec = logical_to_mesh_axes((KV_BATCH, None, None, None), mesh=self.mesh, rules=logical_rules)
qkv_pspec = remove_incompatible_mesh_axes_from_partition_spec(
qkv_pspec, qkv.shape, self.mesh, dims=(0,), allow_remove_axes=True
)
b_a_pspec = remove_incompatible_mesh_axes_from_partition_spec(
b_a_pspec, b.shape, self.mesh, dims=(0,), allow_remove_axes=True
)
conv_state_pspec = remove_incompatible_mesh_axes_from_partition_spec(
conv_state_pspec, conv_state_arg.shape, self.mesh, dims=(0,), allow_remove_axes=True
)
recurrent_state_pspec = remove_incompatible_mesh_axes_from_partition_spec(
recurrent_state_pspec, recurrent_state_arg.shape, self.mesh, dims=(0,), allow_remove_axes=True
)
if cfg.shard_mode == ShardMode.EXPLICIT:
qkv = jax.sharding.reshard(qkv, jax.sharding.NamedSharding(self.mesh, qkv_pspec))
b = jax.sharding.reshard(b, jax.sharding.NamedSharding(self.mesh, b_a_pspec))
a = jax.sharding.reshard(a, jax.sharding.NamedSharding(self.mesh, b_a_pspec))
conv_state_arg = jax.sharding.reshard(
conv_state_arg, jax.sharding.NamedSharding(self.mesh, conv_state_pspec)
)
recurrent_state_arg = jax.sharding.reshard(
recurrent_state_arg, jax.sharding.NamedSharding(self.mesh, recurrent_state_pspec)
)| decay_dtype = getattr(cfg, "gdn_decay_dtype", jnp.float32) | ||
| if isinstance(decay_dtype, str): | ||
| decay_dtype = getattr(jnp, decay_dtype, jnp.float32) | ||
| state_dtype = getattr(cfg, "gdn_state_dtype", jnp.float32) | ||
| if isinstance(state_dtype, str): | ||
| state_dtype = getattr(jnp, state_dtype, jnp.float32) |
There was a problem hiding this comment.
The manual type checks and getattr calls on jnp can be simplified and made more robust by using jnp.dtype directly. jnp.dtype natively handles both string representations (e.g., 'float32') and existing dtype objects.
decay_dtype = jnp.dtype(getattr(cfg, "gdn_decay_dtype", jnp.float32))
state_dtype = jnp.dtype(getattr(cfg, "gdn_state_dtype", jnp.float32))
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Rohan-Bierneni
changed the title
b5061ab to
acdf62c
Compare
… in Qwen3 model - Wire gdn_decoupled_conv1d into Qwen3NextGatedDeltaNet in qwen3.py, gated by use_gdn_kernel configuration flag. - Add gdn_state_dtype: "float32" and gdn_decay_dtype: "float32" configurations to base.yml, types.py, and pyconfig.py for hybrid precision stability. - Retain pure JAX triangular solve optimization when use_gdn_kernel is disabled. - Add pyconfig validation tests for GDN configuration parameters.
acdf62c to
a8d6db0
Compare
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 4 of 5 in the stacked series enabling the Pallas Gated Delta Net (GDN) backward pass kernel in MaxText.
This PR integrates the GDN backward pass kernel into the Qwen3 model architecture with full configuration controls:
src/maxtext/models/qwen3.py):gdn_decoupled_conv1dintoQwen3NextGatedDeltaNet, cleanly toggled bycfg.use_gdn_kernel.use_gdn_kernelis False.use_gdn_kernel(default:False),gdn_state_dtype: "float32", andgdn_decay_dtype: "float32"tosrc/maxtext/configs/base.yml,types.py, andpyconfig.py.float32while keeping activations inbfloat16.tests/unit/pyconfig_test.pywith validation coverage for GDN configuration knobs.Files Changed
src/maxtext/models/qwen3.py: Wiringgdn_decoupled_conv1dintoQwen3NextGatedDeltaNet.src/maxtext/configs/base.yml: Default settings foruse_gdn_kernel,gdn_state_dtype, andgdn_decay_dtype.src/maxtext/configs/types.py: Type annotations and schema definitions for GDN fields.src/maxtext/configs/pyconfig.py: Parser validation for GDN flags.tests/unit/pyconfig_test.py: Unit test verifying pyconfig parsing and validation for GDN knobs.Tests
pre-commit run --files ...(all hooks passed:codespell,pylint,pyink).pytest tests/unit/pyconfig_test.py(41/41 unit tests passed).Checklist