[NNX] Delete Linen 3/5: collapse dispatch in inference (maxengine, kvcache, vLLM, LoRA) by ecnal-cienet · Pull Request #4359 · AI-Hypercomputer/maxtext

ecnal-cienet · 2026-07-06T14:52:53Z

NNX Migration Route Map

✅ 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)
✅ 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)
✅ 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)
✅ 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)
✅ NNX correctness fixes, feature enablements, and vocab tiling on NNX.
✅ NNX-native DPO.
✅ NNX-native MaxEngine inference. (PR [NNX] NNX migration prep (7/N): NNX-native MaxEngine inference #3821)
✅ NNX-native LoRA + GRPO. (PR [NNX] NNX migration prep (8/N): NNX native lora grpo #3824)
✅ NNX-aware QK-Clip + remaining 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)
✅ Vocab tiling custom_vjp for NNX.
✅ Set NNX defaults to True; regenerate sharding goldens; flip back integration-test pure_nnx=False annotations. (PR [NNX] NNX migration (11/N): set pure_nnx / enable_nnx / pure_nnx_decoder defaults to True #3526)
Delete Linen-specific code paths and NNX compatibility flags — delivered as 5 stacked PRs (12.1–12.5):
12.1. ✅ Collapse Linen dispatch in utils / quantization / checkpointing / sharding. (PR [NNX] Delete Linen 1/5: collapse dispatch in utils, quantization, checkpointing, sharding #4357)
12.2. ✅ Collapse Linen dispatch in trainers (pre-train, DiLoCo, GRPO); regenerate the NNX DPO golden. (PR [NNX] Delete Linen 2/5: collapse dispatch in trainers (pre-train, DiLoCo, GRPO) #4358)
12.3. 🔄 [This PR] Collapse Linen dispatch in inference / serving — maxengine.py, kvcache.py, vllm_decode.py, lora_utils.py + tests. Also drops the enable_nnx / pure_nnx_decoder entries vllm_decode.py forwarded into the nested vLLM maxtext_config. Stacks on 12.2.
12.4. ❌ Delete the Linen decoder/attention/embedding layers + *_as_linen model wrappers.
12.5. ❌ Remove the pure_nnx / enable_nnx / pure_nnx_decoder config flags.

Description

Third of the stacked "Delete Linen" series (see 1/5). Collapses the Linen dispatch in the inference and serving path: maxengine.py, kvcache.py, vllm_decode.py, and lora_utils.py, with their tests updated. This also drops the enable_nnx / pure_nnx_decoder entries that vllm_decode.py forwarded into the nested vLLM maxtext_config, since that nested MaxText now defaults to NNX.

Serving already ran the NNX path by default, so this is dead-branch removal with no behavior change. Base branch is feat/nnx-del-linen-2-trainers. ~8 files, +189 / −648.

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 · 2026-07-06T15:02:43Z

Codecov Report

❌ Patch coverage is 71.16279% with 62 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/trainers/pre_train/train.py	69.23%	31 Missing and 13 partials ⚠️
src/maxtext/utils/lora_utils.py	0.00%	12 Missing ⚠️
src/maxtext/utils/train_utils.py	83.33%	2 Missing and 1 partial ⚠️
src/maxtext/trainers/pre_train/train_compile.py	87.50%	1 Missing and 1 partial ⚠️
src/maxtext/utils/muon_utils.py	90.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

…ate setup and train loop The pure_nnx defaults are true, so the Linen branches in the pre-train state and train-loop path are dead. Collapse them: - train_utils.setup_train_loop: always build the abstract NNX model and a TrainStateNNX init_state_fn; drop the Linen model/TrainState branch and the Linen arms of the DiLoCo sharding and debug_sharding blocks. - maxtext_utils: get_functional_train_with_signature and get_functional_eval_with_signature drop the trailing rng in_sharding; load_compiled drops the example rng; get_abstract_state delegates to get_abstract_state_nnx. setup_initial_state is deliberately left alone. Its Linen branch is entangled with the checkpoint restore overlay, and the orbax v1 migration is touching that code; it is collapsed in a later change. Tests follow the same narrowing: the Linen-only cases in maxtext_utils_test, state_dtypes_test and the sharding_compare_test Linen-golden driver go away. test_deepseek4 is skipped rather than pinned to Linen, since nnx_decoders.py has no deepseek4 decoder_block branch yet.

… muon utils - sharding.maybe_update_params_sharding_with_opt delegates to the _nnx variant; build_zero1_input_state_mesh_shardings drops its Linen early return. - muon_utils.get_muon_weight_dimension_numbers drops the isinstance(nnx.Module) test and the Linen get_abstract_param path; get_model_mdn loses its pure_nnx parameter and always builds the abstract NNX model. - run_sharding_dump loses its --pure_nnx flag. Tests follow: the Linen branch test in muon_utils_test goes away, optimizers_test drops the pure_nnx argument and the dual-shape comparison, sharding_nnx_test drops the flag from its fake config.

… and model creation - quantizations.maybe_quantize_model always runs the qwix forward pass with the dummy tokens/positions/segment ids (and the MTP decoder targets when mtp_num_layers > 0), then pops the transient nnx.Intermediate variables the traced forward sows. - model_creation_utils.from_pretrained always builds the sharded model through maxtext_utils_nnx.create_nnx_sharded_model. Tests follow: quantizations_test and nnx_quant_guard_test drop their flag arguments and Linen expectations, correctness_tests_nnx_dispatch_test keeps only the NNX case, and forward_pass_logit_checker always loads via from_pretrained.

…iLoCo, GRPO)

…kvcache, vLLM, LoRA)

This was referenced Jul 6, 2026

[NNX] Delete Linen 4/5: remove the Linen decoder/attention layers and *_as_linen model wrappers #4360

Draft

[NNX] Delete Linen 5/5: remove the pure_nnx / enable_nnx / pure_nnx_decoder config flags #4361

Draft

ecnal-cienet force-pushed the feat/nnx-del-linen-3-inference branch 2 times, most recently from cb9c0d0 to fbe1484 Compare July 7, 2026 21:00

ecnal-cienet force-pushed the feat/nnx-del-linen-3-inference branch 7 times, most recently from b514961 to 6a2b270 Compare July 21, 2026 18:55

ecnal-cienet force-pushed the feat/nnx-del-linen-3-inference branch 6 times, most recently from 2a26f98 to e1a43f7 Compare July 29, 2026 14:48

ecnal-cienet force-pushed the feat/nnx-del-linen-3-inference branch 3 times, most recently from 24c27a8 to dab8d49 Compare August 19, 2026 16:46

ecnal-cienet added 5 commits August 20, 2026 21:29

[NNX] Delete Linen (2/5): collapse dispatch in trainers (pre-train, D…

34e33cd

…iLoCo, GRPO)

[NNX] Delete Linen (3/5): collapse dispatch in inference (maxengine, …

5890a05

…kvcache, vLLM, LoRA)

ecnal-cienet force-pushed the feat/nnx-del-linen-3-inference branch from dab8d49 to 5890a05 Compare August 20, 2026 21:38

Navigation Menu

Uh oh!

[NNX] Delete Linen 3/5: collapse dispatch in inference (maxengine, kvcache, vLLM, LoRA) - #4359

[NNX] Delete Linen 3/5: collapse dispatch in inference (maxengine, kvcache, vLLM, LoRA)#4359
ecnal-cienet wants to merge 5 commits into
mainfrom
feat/nnx-del-linen-3-inference

ecnal-cienet commented Jul 6, 2026

Uh oh!

codecov Bot commented Jul 6, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

ecnal-cienet commented Jul 6, 2026

NNX Migration Route Map

Description

Checklist

Uh oh!

codecov Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

codecov Bot commented Jul 6, 2026 •

edited

Loading