Mhc expand - #5014
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.
![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 adds support for Qwen3-Next models, including new configurations, scannable blocks, and a hybrid Gated Delta Net (GDN) implementation. It also integrates Manifold-Constrained Hyper Connections (mHC) and updates the Muon optimizer dimension mapping to be invariant to layer scanning. Critical feedback identifies several issues: an AttributeError in parameter mapping, a memory regression (OOM risk) in nnx_decoders.py from returning the full nnx.state inside a scan, a dangerous bare except: pass block, potential shape mismatches in Tokamax group sizes when inputs are not perfectly divisible, and a crash on non-TPU platforms due to the removal of a platform check in ragged_gather.py.
| for i in range(config.base_num_decoder_layers): | ||
| prefix = f"params-decoder-layers_{i}" | ||
| block_idx = i % config.inhomogeneous_layer_cycle_interval | ||
| is_full_attention_layer = ( | ||
| block_idx + 1 | ||
| ) % config.inhomogeneous_layer_cycle_interval == 0 | ||
| _attach_block_hooks(prefix, is_global=is_full_attention_layer) |
There was a problem hiding this comment.
Using config.base_num_decoder_layers and config.inhomogeneous_layer_cycle_interval will raise an AttributeError because config is the HuggingFace configuration dictionary (or object), which does not contain these MaxText-specific fields. Instead, you should use the local variables num_main_layers and layer_cycle_interval which are already defined in the outer scope.
| for i in range(config.base_num_decoder_layers): | |
| prefix = f"params-decoder-layers_{i}" | |
| block_idx = i % config.inhomogeneous_layer_cycle_interval | |
| is_full_attention_layer = ( | |
| block_idx + 1 | |
| ) % config.inhomogeneous_layer_cycle_interval == 0 | |
| _attach_block_hooks(prefix, is_global=is_full_attention_layer) | |
| for i in range(num_main_layers): | |
| prefix = f'params-decoder-layers_{i}' | |
| block_idx = i % layer_cycle_interval | |
| is_full_attention_layer = ( | |
| block_idx + 1 | |
| ) % layer_cycle_interval == 0 | |
| _attach_block_hooks(prefix, is_global=is_full_attention_layer) |
| # 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 the full nnx.state(layer) (which includes read-only nnx.Param variables) inside the scan body will cause JAX to stack all parameters along the scan dimension. This re-introduces the huge unnecessary memory allocation that the original code explicitly avoided to prevent Out-Of-Memory (OOM) issues during compilation and execution. Please revert to splitting the state and returning only the updated non-parameter state.
| return new_carry, nnx.state(layer) | |
| _, _, updated_state = nnx.split(layer, nnx.Param, ...) | |
| return new_carry, updated_state |
| try: | ||
| forked_rngs = rngs.fork(split=length) | ||
| except: # pylint: disable=bare-except | ||
| pass |
There was a problem hiding this comment.
The bare except: pass block is highly dangerous here. If rngs.fork raises an exception, forked_rngs will remain undefined, which will immediately cause a NameError or UnboundLocalError on the next line when nnx.split(forked_rngs) is called. Since rngs is a required argument and must be valid, this try-except block should be removed.
forked_rngs = rngs.fork(split=length)| 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 inputs.shape[0] is perfectly divisible by num_groups. If there is a remainder, the sum of the group sizes will be less than the total number of inputs, which can cause shape mismatch or out-of-bounds errors in Tokamax. Please distribute the remainder evenly across the groups.
| 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, | |
| ) |
| if enforce_fallback: | ||
| return _fallback_implementation(x, indices, weights, has_weights) |
There was a problem hiding this comment.
Removing the TPU platform check or jax.devices()[0].platform != "tpu" means that running this function on non-TPU platforms (such as CPU during unit tests or local development) will attempt to call pltpu.get_tpu_info(), which will crash. Please restore the platform check to safely guard against non-TPU execution.
| if enforce_fallback: | |
| return _fallback_implementation(x, indices, weights, has_weights) | |
| if enforce_fallback or jax.devices()[0].platform != 'tpu': | |
| return _fallback_implementation(x, indices, weights, has_weights) |
bd34611 to
5db6100
Compare
5db6100 to
c4ad690
Compare
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.