Fix NNX MTP zero-loss, qwix MTP crash, and DeepSeek scan-axis sharding corruption

[ecnal-cienet]

Fixes three bugs in the NNX path surfaced by running DeepSeek-V3 with multi-token prediction. They are independent; the third is not an MTP bug at all and was hiding behind the second.

1. mtp_loss is silently 0 under NNX (b/535659954)

mtp_losses / mtp_acceptance subclass nnx.Intermediate, and NNX type filters match by isinstance. So the generic nnx.pop(model, nnx.Intermediate) in loss_fn consumed them before the dedicated pop ran, which then returned {}; calculate_mtp_loss fell through to its return 0.0 default. No error, just a wrong number.

The comment claiming these variables are "not Intermediate, so the nnx.pop above misses them" was the false premise behind the bug.

Fix: pop the specific subclasses before the generic nnx.Intermediate pop. Linen is unaffected — it keys mutable collections by literal name.

2. TypeError: roll requires ndarray ... got NoneType with qwix (b/535662227)

NNX modules are stateful, so qwix.quantize_model must run a real forward pass. MaxText passed only tokens/positions/segment_ids, so the MTP block received target_ids=None and jnp.roll raised. Linen never hits this because quantize_linen_model rejects model inputs and never traces a forward.

Fix: supply dummy targets at that call site, gated on mtp_num_layers > 0, restoring the semantics of the is_initializing() guard that the NNX port left commented out in models.py.

3. DeepSeek params silently replicated under NNX (pre-existing)

Surfaced once (2) was fixed: AssertionError: Unsharded parameter percentage (2.60%) exceeds tolerance (2.00%).

_create_scanned_layers records each stack's axis name in nnx.PARTITION_NAME — dense_layers / moe_layers for DeepSeek — but _apply_layers_sequentially hardcoded the literal "layers" when round-tripping the scan axis. The name never matched, so the fallback while len(l) > ndim: l.pop(0) deleted the real leading logical axis and nnx_add_scan_axis inserted a bogus 'layers':

decoder/dense_layers/mlp/wi_0/kernel
  before: ('embed', 'dense_layers', 'mlp')
  after : ('dense_layers', 'layers', 'mlp')   <- after ONE forward pass

'embed' is the only carrier of fsdp, so the params became P(None, ...).

This is neither an MTP bug nor a qwix bug: it reproduces at mtp_num_layers=0 and with qwix entirely absent (29 of 32 params corrupted by a single plain forward). It normally stays invisible because get_abstract_state_nnx snapshots shardings before the first train step, so the damage lands on a discarded traced copy. qwix makes it fatal only because it quantizes from inside model construction, i.e. inside the traced snapshot. Standard models pass metadata_axis_name="layers", matching the literal — only DeepSeek-style multi-stack models are exposed.

Fix: resolve the axis name per-variable from nnx.PARTITION_NAME rather than a caller literal, and raise instead of blindly stripping an axis when the metadata is inconsistent. Flax's own nnx.spmd.remove_axis raises on this same mismatch; this reimplementation corrupted silently.

sharding_tolerance is untouched — the params are genuinely sharded again.

Test

python3 -m maxtext.trainers.pre_train.train src/maxtext/configs/base.yml \
  base_output_directory=/home/wanglance_google_com/tmp/out run_name=nnx_mtp \
  model_name=deepseek3-tiny per_device_batch_size=1 max_target_length=128 steps=20 \
  dataset_type=synthetic enable_checkpointing=false sparse_matmul=false \
  mtp_num_layers=1 mtp_loss_scaling_factor=0.1

• enable_nnx=false pure_nnx_decoder=false pure_nnx=false for Linen, + use_qwix_quantization=true quantization=fp8_full for qwix.

• nnx_mtp
• nnx_mtp_qwix
• linen_mtp
• linen_mtp_qwix

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]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 89.28571% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/utils/maxtext_utils_nnx.py	80.00%	1 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.