Dipakg custom qwen3 0825 - #5010
Conversation
…uon sharding, and mHC-lite Pallas/Mosaic kernels on TPU v6e-256. - Integrate DeepSeek mHC-lite architecture and high-performance Mosaic/Pallas TPU kernels with custom VJP rules and SPMD shard_map integration. - Support Qwen3-Next 80B architecture with Gated Delta Net (GDNv3) and Pallas custom VJP in hybrid_gdn.py. - Configure Muon optimizer and sharding rules for MoE parameters (shard-exp-on-fsdp). - Update AOT compile, XManager Borg, and XPK launch scripts for Qwen3-Next 80B on TPU v6e-256 (glp_16x16) topology to enable mHC-lite and Pallas kernel flags. - Add and update unit tests for mHC, GDN, Muon utilities, NNX decoders, and parameter mapping. - Validated clean AOT compilation on TPU v6e-256 via qn80b-v6e256-compile.sh and unit tests via mhc_test.
- Integrates sharding-aware Muon optimizer supporting Newton-Schulz iterations and partition padding into MaxText. - Fixes weight dimension indexing in muon_utils.py to account for layer scanning, 1D tensor exclusions, and rank-aware 2D attention projection handling. - Fixes routing gate exclusion logic to use exact segment matching across gate and router variants. - Plumbs NamedSharding in sharding.py by collecting momentum trees across multi-branch partitioned optimizers, extracting MuonState NamedTuples, and filtering optax.MaskedNode placeholders. - Threads mesh through train_utils.py, train_compile.py, and maxtext_engine.py to get_optimizer. - Add some minor tests throughout.
…aluation. - Replaced jnp.where boolean masking with float multiplication inside the chunk rules. - Moved sequence padding and query scaling logic before multi-head expansion to eliminate redundant FLOPs. - Removed explicit jnp.repeat calls by leveraging implicit broadcasting via jnp.einsum. - Fused transposes and replaced chained jnp.matmul operations with optimized jnp.einsum calls. - Configured DenseGeneral projections to natively handle multi-dimensional shapes and replaced jnp.split with array slicing. Performance increased on TPU v6e-256 from 2600Tokens/s/device) to ~2640 Tokens/s/device which is about 1.5% improvement.
![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 support for the Qwen3-Next model architecture, including specialized Gated Delta Net (GDN) implementations and MoE configurations. It updates the Muon optimizer to support all-to-all communication and configurable Newton-Schulz steps, and refactors the NNXDecoder to support scanned blocks for Qwen3-Next. The reviewer identified several critical issues, including potential crashes due to missing platform checks in ragged gather kernels, incorrect transpose permutations in decoder blocks, key collision bugs in bias update path construction, and performance regressions in memory allocation during scanning. Additionally, the reviewer noted a potential logic error in group size distribution for ragged dot operations.
| # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests). | ||
| # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU. | ||
| if enforce_fallback or jax.devices()[0].platform != "tpu": | ||
| if enforce_fallback: |
There was a problem hiding this comment.
Removing the platform check or jax.devices()[0].platform != "tpu" will cause pltpu.get_tpu_info() to be executed on non-TPU platforms (such as CPU during unit tests), which will result in a crash. Please restore the platform check to safely fall back on non-TPU hardware.
| if enforce_fallback: | |
| if enforce_fallback or jax.devices()[0].platform != "tpu": |
| # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests). | ||
| # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU. | ||
| if enforce_fallback or jax.devices()[0].platform != "tpu": | ||
| if enforce_fallback: |
There was a problem hiding this comment.
Removing the platform check or jax.devices()[0].platform != "tpu" will cause pltpu.get_tpu_info() to be executed on non-TPU platforms (such as CPU during unit tests), which will result in a crash. Please restore the platform check to safely fall back on non-TPU hardware.
| if enforce_fallback: | |
| if enforce_fallback or jax.devices()[0].platform != "tpu": |
| # ========================================================================= | ||
| o = o_chunks.transpose(1, 0, 3, 2, 4) | ||
| o = o.reshape(B, -1, H, V_dim) | ||
| o = o_chunks.transpose(0, 1, 4, 2, 3, 5).reshape(B, seq_len, H_k * R, V_dim) |
There was a problem hiding this comment.
The transpose permutation (0, 1, 4, 2, 3, 5) keeps num_chunks at axis 0 and B (batch size) at axis 1. When this is reshaped to (B, seq_len, H_k * R, V_dim), it will scramble the batch dimension because B is not the leading dimension before the reshape. To keep the batch dimension intact, the batch axis B must be transposed to axis 0 first using the permutation (1, 0, 4, 2, 3, 5).
| o = o_chunks.transpose(0, 1, 4, 2, 3, 5).reshape(B, seq_len, H_k * R, V_dim) | |
| o = o_chunks.transpose(1, 0, 4, 2, 3, 5).reshape(B, seq_len, H_k * R, V_dim) |
| moe_bias_updates = { | ||
| "/".join(tuple(k.key for k in path if hasattr(k, "key"))): val | ||
| for path, val in jax.tree_util.tree_leaves_with_path( | ||
| intermediate_outputs | ||
| ) | ||
| if tuple(k.key for k in path if hasattr(k, "key"))[-1:] | ||
| == ("moe_bias_updates",) | ||
| } |
There was a problem hiding this comment.
Using tuple(k.key for k in path if hasattr(k, "key")) to construct the path string skips all non-DictKey components, such as SequenceKey (which has idx instead of key). For sequential/non-scanned layers, this causes all layers' moe_bias_updates to map to the exact same key (e.g., "decoder/layers/moe_bias_updates"), clobbering previous layers' updates and keeping only the last one. Please include the sequence index idx in the path string to prevent key collisions.
| moe_bias_updates = { | |
| "/".join(tuple(k.key for k in path if hasattr(k, "key"))): val | |
| for path, val in jax.tree_util.tree_leaves_with_path( | |
| intermediate_outputs | |
| ) | |
| if tuple(k.key for k in path if hasattr(k, "key"))[-1:] | |
| == ("moe_bias_updates",) | |
| } | |
| moe_bias_updates = { | |
| "/".join( | |
| tuple( | |
| str(k.key) if hasattr(k, "key") | |
| else (str(k.idx) if hasattr(k, "idx") else str(k)) | |
| for k in path | |
| ) | |
| ): val | |
| for path, val in jax.tree_util.tree_leaves_with_path( | |
| intermediate_outputs | |
| ) | |
| if len(path) > 0 and ( | |
| hasattr(path[-1], "key") and path[-1].key == "moe_bias_updates" | |
| ) | |
| } |
| param_paths = [ | ||
| tuple(k.key for k in path if hasattr(k, "key")) | ||
| for path, _ in jax.tree_util.tree_leaves_with_path( | ||
| nnx.state(model, nnx.Param).to_pure_dict() | ||
| ) | ||
| ] |
There was a problem hiding this comment.
Using tuple(k.key for k in path if hasattr(k, "key")) to construct the parameter paths skips all non-DictKey components, such as SequenceKey (which has idx instead of key). This causes a mismatch with the moe_bias_updates keys constructed in train.py for sequential/non-scanned layers, leading to a ValueError during training. Please include the sequence index idx in the path tuple to ensure correct matching.
param_paths = [
tuple(
str(k.key) if hasattr(k, "key")
else (str(k.idx) if hasattr(k, "idx") else str(k))
for k in path
)
for path, _ in jax.tree_util.tree_leaves_with_path(
nnx.state(model, nnx.Param).to_pure_dict()
)
]| # This prevents huge unnecessary memory allocation. | ||
| _, _, updated_state = nnx.split(layer, nnx.Param, ...) | ||
| return new_carry, updated_state | ||
| return new_carry, nnx.state(layer) |
There was a problem hiding this comment.
Returning nnx.state(layer) directly from the scan body instead of splitting out nnx.Param will cause JAX to stack all read-only parameters along the scan dimension, leading to huge unnecessary memory allocations during the forward pass. Please restore the nnx.split logic to avoid returning and stacking read-only parameters.
| return new_carry, nnx.state(layer) | |
| _, _, updated_state = nnx.split(layer, nnx.Param, ...) | |
| return new_carry, updated_state |
| num_groups = group_sizes.shape[0] | ||
| avg_size = inputs.shape[0] // num_groups | ||
| return tokamax.RaggedDotGroupSizes( | ||
| group_sizes, | ||
| inputs.shape[0], | ||
| (avg_size,) * num_groups, | ||
| ) |
There was a problem hiding this comment.
Using (avg_size,) * num_groups assumes that the total number of tokens inputs.shape[0] is perfectly divisible by num_groups. If it is not, the sum of the group sizes in the returned RaggedDotGroupSizes will not match the total number of tokens, which can lead to correctness or runtime errors in tokamax. Consider distributing the remainder across the groups to ensure the sum matches inputs.shape[0] exactly.
| num_groups = group_sizes.shape[0] | |
| avg_size = inputs.shape[0] // num_groups | |
| return tokamax.RaggedDotGroupSizes( | |
| group_sizes, | |
| inputs.shape[0], | |
| (avg_size,) * num_groups, | |
| ) | |
| num_groups = group_sizes.shape[0] | |
| avg_size = inputs.shape[0] // num_groups | |
| remainder = inputs.shape[0] % num_groups | |
| sizes = tuple(avg_size + (1 if i < remainder else 0) for i in range(num_groups)) | |
| return tokamax.RaggedDotGroupSizes( | |
| group_sizes, | |
| sizes, | |
| ) |
2 participants
Description
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
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.