Fix NaN loss by using l2Norm in GDN v3 kernel

[muskansh-google]

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:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

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):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces support for the Qwen3-Next-80B model in MaxText, including architectural updates, hierarchical nested scans for inhomogeneous layers, and a hybrid Gated Delta Net (GDN) implementation combining Tokamax forward and Pallas custom backward passes. It also adds several run scripts, updates parameter mapping, integrates Manifold-Constrained Hyper Connections (mHC), and refines Muon optimizer compatibility. The reviewer feedback highlights critical issues that need to be addressed: explicitly loading Pallas VMEM Ref buffers using [...] in the hybrid GDN backward kernel, using nnx_sync_moveaxis instead of jnp.moveaxis to keep sharding metadata synchronized in nnx_decoders.py and qwen3.py, removing a bare except: pass block in nnx_decoders.py to prevent potential NameErrors, and tightening the type check in sharding.py to prevent TypeErrors on non-dict containers.

[gemini-code-assist]

These VMEM buffers are Ref objects in Pallas. Slicing them with [slot, ...] returns another Ref view, not a JAX array. To load the values into registers as JAX arrays so they can be used in JAX operations (like .reshape, .astype, or jax.vjp), you must explicitly load them using [...] or .load().

Suggested change

	padded_pre_conv_qkv_val = padded_pre_conv_qkv_vmem[slot, ...]
	qkv_val = qkv_vmem[slot, ...]
	b_val = b_vmem[slot, ...]
	a_val = a_vmem[slot, ...]
	do_val = do_vmem[slot, ...]
	state_prev_val = chunk_states_vmem[slot, ...]
	padded_pre_conv_qkv_val = padded_pre_conv_qkv_vmem[slot, ...][...]
	qkv_val = qkv_vmem[slot, ...][...]
	b_val = b_vmem[slot, ...][...]
	a_val = a_vmem[slot, ...][...]
	do_val = do_vmem[slot, ...][...]
	state_prev_val = chunk_states_vmem[slot, ...][...]

[gemini-code-assist]

Using jax.tree.map with jnp.moveaxis directly on scanned_params only moves the axis of the array values, leaving the sharding metadata (such as NamedSharding or PartitionSpec) out of sync. Use the newly added maxtext_utils_nnx.nnx_sync_moveaxis helper function to keep both the values and sharding metadata synchronized.

Suggested change

	scanned_params = jax.tree.map(lambda x: jnp.moveaxis(x, 0, scan_axis), scanned_params)
	scanned_params = maxtext_utils_nnx.nnx_sync_moveaxis(scanned_params, 0, scan_axis)

[gemini-code-assist]

Import maxtext_utils_nnx alongside maxtext_utils to support calling nnx_sync_moveaxis in the scannable block.

Suggested change

	from maxtext.utils import maxtext_utils
	from maxtext.utils import maxtext_utils, maxtext_utils_nnx

[gemini-code-assist]

Directly calling jnp.moveaxis on stacked_params leaves the sharding metadata out of sync with the array values. Use maxtext_utils_nnx.nnx_sync_moveaxis to ensure both the values and sharding metadata are updated correctly.

Suggested change

	stacked_state = nnx.State.merge(stacked_params, stacked_other)
	stacked_params = maxtext_utils_nnx.nnx_sync_moveaxis(stacked_params, 0, scan_axis)

[gemini-code-assist]

A bare except: pass block is a code smell that can silence unexpected errors. Furthermore, if rngs.fork fails, forked_rngs will be left undefined, leading to a NameError on the very next line. Since rngs is a required argument of type nnx.Rngs, the try...except block should be removed entirely.

    forked_rngs = rngs.fork(split=length)

[gemini-code-assist]

Checking hasattr(obj, "__getitem__") is too broad and can match non-dict containers like strings (where "mu" in obj would be true but indexing with a string key like obj["mu"] will raise a TypeError). Since the optimizer moment state is expected to be a dictionary or an nnx.State, it is safer to explicitly check for these types.

Suggested change

	if hasattr(obj, "__getitem__") and "mu" in obj:
	if isinstance(obj, (dict, nnx.State)) and "mu" in obj: