[NNX] NNX migration prep (10/N): vocab tiling custom_vjp with output-head carve-out

[ecnal-cienet]

NNX Migration Route Map

1. ✅ Add NNX scaffolding: pure_nnx flag, init_state_fn, TrainStateNNX, NNX utils. Linen workflow unchanged. (PR NNX migration prep (1/N): pure_nnx flag and init_state_fn scaffolding #3427)
2. ✅ NNX sharding utilities: get_abstract_state_nnx, get_named_sharding_nnx, set_named_sharding_nnx, get_partition_spec_nnx, get_mesh_from_config. (PR NNX migration prep (2/N): NNX utils and sharding utilities #3470)
3. ✅ NNX fully supported end-to-end: TrainStateNNX, model creation, gradient accumulation, checkpointing, and training loop dispatch. (PR NNX migration prep (3/N): TrainState, model creation, and end-to-end training loop #3500)
4. ✅ Sharding diagnostics on NNX, plus post-training bugfixes that surfaced once the NNX path got exercised end-to-end. (PR [NNX] NNX migration prep (4/N): sharding tools and post-training fixes #3652)
   4.5. ✅ Linen↔NNX checkpoint converter. (PR [NNX] NNX migration prep (4.5/N): Linen<->NNX checkpoint converter #3843)
   4.6. ✅ Linen↔NNX checkpoint comparator. (PR [NNX] NNX migration prep (4.6/N): Linen<->NNX checkpoint comparator #3846)
5. ✅ NNX correctness fixes, feature enablements, and vocab tiling MVP on NNX.
6. ✅ NNX-native DPO.
7. ✅ NNX-native MaxEngine inference. (PR [NNX] NNX migration prep (7/N): NNX-native MaxEngine inference #3821)
8. ✅ NNX-native LoRA + GRPO. (PR [NNX] NNX migration prep (8/N): NNX native lora grpo #3824)
9. ✅ NNX-aware QK-Clip + NNX-format checkpoint utilities. (PR [NNX] NNX migration prep (9/N): NNX-aware QK-Clip + checkpoint utilities #3836)
   9.5. ✅ NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix. (PR [NNX] NNX migration prep (9.5/N): NNX + AQT in MaxEngine + serve-mode reload + gpt3 prefill fix #3844)
10. 🔄 [This PR] Vocab tiling custom_vjp for NNX (with output-head carve-out, originally scoped as a follow-up PR10.5 — bundled in here).
11. ❌ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations.
12. ❌ Delete Linen-specific code paths and NNX compatibility flags.

Description

Replaces the PR5 NNX vocab-tiling MVP (chunked forward + default autograd backward) with a jax.custom_vjp that mirrors the Linen path's backward-memory savings, then carves out the output-head params so the custom_vjp's residuals + grad accumulator scale with LM-head size, not with the full model. PR10.5 was originally scoped as a separate follow-up but is bundled here — the carve-out and its supporting fixes (nnx_decoders.py .value modernization, explicit zero cotangents in the bwd) are necessary for the NNX path to AOT-compile correctly under num_vocab_tiling > 1 + pure_nnx=True.

The Linen vocab_tiling_linen_loss is byte-for-byte unchanged. Call sites in train.py, pyconfig_deprecated.py, and configs/types.py are unchanged.

Diff: +517 / −44 across 5 files (3 src, 2 test).

What it does

src/maxtext/utils/vocabulary_tiling.py — vocab_tiling_nnx_loss rewrite:

• 3-way nnx.split — outside the custom_vjp, a callable path filter (_is_output_head_param_path) matching {token_embedder, shared_embedding, decoder_norm, logits_dense} separates head_params (the only nnx.Param paths apply_output_head touches) from other_params (transformer layers, etc.) and rest (rngs).
• Custom_vjp primals — (head_params, other_params, rest, hidden_states, labels, segmentation). Only head_params and hidden_states are differentiated; other_params + rest are threaded through as non-differentiated primals so their tracers don't have to cross both the custom_vjp and the inner lax.scan boundary (which previously caused UnexpectedTracerError under logits_via_embedding=True).
• Forward — reshape to (num_vocab_tiling, vocab_tile_size, ...), run lax.scan whose body rebuilds the model per chunk via nnx.merge(graphdef, chunk_head, chunk_other, chunk_rest, copy=True) and calls logits_from_hidden_states. Initial scan accumulator is fp32 (was hidden_states.dtype previously — caused a lax.scan carry dtype mismatch with bf16 hidden_states since cross_entropy_with_logits always returns fp32). Residuals are (chunk_head, chunk_other, chunk_rest, reshaped_*, batch/seq/emb).
• Backward — second lax.scan whose body builds loss_fn_for_vjp = lambda p, h: ..., calls jax.vjp(loss_fn_for_vjp, chunk_head_params, hidden_chunk), accumulates grad_head via tree.map(add), emits per-chunk grad_hidden. Chain-rules grad_head *= loss_cotangent and dtype-casts each grad back to its primal's dtype (custom_vjp requires this). chunk_other_params and chunk_rest cotangents are explicit tree_map(jnp.zeros_like, ...) zero pytrees, NOT None — None makes JAX synthesize zeros at AOT trace time with axis-0 stacking (jax.scan convention) for nnx.scan-stacked transformer-layer params, which carry axis-1 stacking (NNX convention), and the cotangent-shape check fails (Expected cotangent type bfloat16[E,M] for primal type bfloat16[E,M], but got bfloat16[L,E,M]). Materializing the zeros ties the cotangent shape to the primal shape exactly.
• Correctness — logits_from_hidden_states provably depends only on head_params; the gradient w.r.t. other_params through this loss is exactly zero. When train.py also calls the full model forward (which produces hidden_states), transformer-layer gradients flow back through grad_hidden_states → outer backward, unaffected by the carve-out.

src/maxtext/layers/nnx_decoders.py — apply_output_head logits_via_embedding=True branch reads embedding_table = shared_embedding.embedding[...] (modern NNX Variable[Array] API) instead of the deprecated .value shim. The .value shim registers the access in NNX mutation tracking, which JAX detects as a tracer leak when the embedding is closure-captured / threaded across the custom_vjp + lax.scan boundaries. The Linen branch (shared_embedding.variables["params"]["embedding"]) is unchanged.

src/maxtext/models/models.py — deletes dead-code self.hidden_states = None and if self.config.num_vocab_tiling > 1: self.hidden_states = hidden_state from the NNX Transformer class. Two lines left over from an early PR5 implementation idea — neither path actually reads model.hidden_states (Linen reads via mutable=["intermediates"]; NNX reads via nnx.pop(model, nnx.Intermediate) from the decoder's sown ("decoder", "hidden_states") intermediate). Without this fix, AOT compile under pure_nnx=True + num_vocab_tiling>1 raised ValueError: Cannot assign data value of type 'LinearizeTracer' to static attribute 'hidden_states' of Pytree type 'Transformer' — would have silently broken any post-PR11 user with vocab tiling on.

Tests

tests/unit/tiling_test.py — new VocabTilingNNXTest class with 9 TPU tests:

• test_nnx_vocab_tiling_non_tied_embedding / _tied_embedding — loss + grad parity vs. full-vocab xent reference for both LM-head modes.
• test_nnx_vocab_tiling_total_z_loss_value_parity — asserts the second tuple element matches the reference (was untested before).
• test_nnx_vocab_tiling_padded_segmentation — half-padded mask; exercises the segmentation != 0 mask branch and asserts padded loss is strictly less than unpadded.
• test_nnx_vocab_tiling_grad_over_hidden_states — argnums=1 differentiation; exercises the custom_vjp's second-primal cotangent path (grad_reshaped_hidden_states); shape + dtype + value parity.
• test_nnx_vocab_tiling_bf16_hidden_states — bf16 inputs with rtol/atol loosened to 5e-2; asserts grad_h.dtype == bf16 (the bwd dtype-cast preserves the primal's dtype). Caught the fp32-accumulator bug.
• test_nnx_vocab_tiling_z_loss_zero — z_loss_multiplier=0; total_z_loss == 0.0 exactly and grad parity holds.
• test_nnx_vocab_tiling_num_vocab_tiling_variants — runs n ∈ {2, 4, 8} and asserts identical loss + grads (catches off-by-one in vocab_tile_size and scan/reshape interactions).
• test_nnx_vocab_tiling_other_params_get_zero_grad — carve-out invariant: asserts every non-head leaf has gradient exactly zero AND at least one head leaf has non-zero gradient (so the test can't trivially pass by zeroing everything). Catches filter bugs (e.g. forgetting that NNX names the embedder token_embedder while Linen names it shared_embedding) and bwd zero-shape bugs.

tests/unit/train_compile_test.py:

• Removed the now-stale pytest.skip("Vocab tiling not supported on NNX.") in test_vocab_tiling_bf16.
• Added test_vocab_tiling_bf16_nnx (@pytest.mark.cpu_only) which AOT-compiles the train step under pure_nnx=true + enable_nnx=true + pure_nnx_decoder=true with num_vocab_tiling=4 and weight_dtype=bfloat16. Surfaced both the models.py dead-code regression and the cotangent-axis-ordering issue the explicit-zeros bwd fixes.

Existing tests untouched.

Stats

• Diff: +517 / −44 across 5 files (0 new, 5 modified).
• Production-code impact: vocab_tiling_nnx_loss rewritten (replaces PR5's MVP). nnx_decoders.py::apply_output_head 1-line modernization affects every NNX run with logits_via_embedding=True. models.py 2-line dead-code deletion has no observable behavior change (no readers).
• Linen preservation: vocab_tiling_linen_loss byte-for-byte unchanged. nnx_decoders.py Linen branch unchanged. TransformerLinenPure unchanged.

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.

[codecov]

Codecov Report

❌ Patch coverage is 29.45892% with 1056 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...xtext/checkpoint_conversion/linen_nnx_converter.py	8.97%	283 Missing and 1 partial ⚠️
src/maxtext/utils/lora_utils.py	5.17%	110 Missing ⚠️
src/maxtext/trainers/pre_train/train.py	43.71%	81 Missing and 22 partials ⚠️
...rc/maxtext/utils/generate_param_only_checkpoint.py	14.03%	96 Missing and 2 partials ⚠️
src/maxtext/utils/maxtext_utils.py	43.33%	60 Missing and 8 partials ⚠️
src/maxtext/utils/layerwise_quantization.py	9.58%	66 Missing ⚠️
src/maxtext/utils/qk_clip_utils.py	10.60%	59 Missing ⚠️
src/maxtext/utils/sharding.py	5.55%	50 Missing and 1 partial ⚠️
src/maxtext/trainers/post_train/dpo/dpo_utils.py	4.00%	48 Missing ⚠️
src/maxtext/utils/model_creation_utils.py	56.52%	25 Missing and 5 partials ⚠️
... and 15 more

📢 Thoughts on this report? Let us know!