Fix NNX sequential layer naming and PyTree structure

[xibinliu]

Description

This is for checkpoint loading

Fix the Errors:

1. DeepSeek Unscanned Checkpoint Mismatch:

When restoring unscanned DeepSeek checkpoints, loading failed with:

ValueError: Checkpoint structure mismatch: 374 of 377 model parameter paths were not found... 
Example missing paths: decoder.dense_layer_0...

Cause: _init_sequential_deepseek registered layers as singular dense_layer and moe_layer, producing PyTree paths like decoder.dense_layer_0. However, Linen checkpoints and MaxText converter scripts format layer keys as plural (decoder.dense_layers_0, decoder.moe_layers_0).

2. Generic Unscanned Model Checkpoint Mismatch (GPT-OSS, Llama4, etc.):

When restoring unscanned generic models, loading failed with:

ValueError: Checkpoint structure mismatch: 456 of 459 model parameter paths were not found
... 
Example missing paths: decoder.layers.0
...

Cause: _init_sequential_generic used _create_and_register_layer which appended to self.layers initialized as nnx.List([]). In Flax NNX, tracking submodules in an nnx.List container forces state PyTree extraction to use list-index paths (decoder.layers.0) instead of named attribute paths (decoder.layers_0), breaking compatibility with on-disk Linen checkpoints.

Summary of Fixes:

• Updated _init_sequential_deepseek to use plural prefixes (dense_layers, moe_layers).
• Used _create_and_register_named_layer in sequential initialization (_init_sequential_deepseek, _init_sequential_generic) so submodules register exclusively under named string attributes (layers_0, layers_1).
• Replaced nnx.List with a plain Python list for self.layers and updated the forward pass to retrieve named attributes via getattr(self, f'layers_{lyr}').
• MaxEngine KV-Cache Stacking & Sharding (maxengine.py): Updated self.prefill_kv_cache_shardings, _maybe_stack_prefill_result_cache, and _maybe_unstack_prefill_result_cache to handle named per-layer keys (layers_0, layers_1) in addition to stacked layers keys when scan_layers=False.
• LoRA Module Regex Pattern Matching (lora_utils.py & lora_utils_test.py):
  - Updated _get_lora_module_path regex replacement to r'layers(?:_[0-9]+|/[0-9]+)?/' so LoRA target verification matches named attributes (decoder/layers_0/...) as well as indexed (decoder/layers/0/...) and scanned (decoder/layers/...) paths.
  - Updated expected regex string assertions in test_get_lora_module_path.

FIXES: b/532719643
FIXES: b/532731449
FIXES: b/532916578

Tests

1. llama4-17b-16e

python3 -m tests.utils.forward_pass_logit_checker \
    src/maxtext/configs/base.yml \
    tokenizer_path=meta-llama/Llama-4-Scout-17B-16E \
    load_parameters_path=gs://ml-auto-solutions/output/jax_models_and_performance/maxtext/chained_tests_llama4_nightly-2026-07-06-01-30-12/unscanned/0/items \
    model_name=llama4-17b-16e \
   ...

2. gpt-oss-20b

python3 -m tests.utils.forward_pass_logit_checker \
    src/maxtext/configs/base.yml \
    model_name=gpt-oss-20b \
    load_parameters_path=gs://ml-auto-solutions/output/unowned/maxtext_stable_gpt-oss-20b-v5p-8-2026-07-06-04-15-19/unscanned/0/items \
   ...

5. deepseek2-16b

python3 -m maxtext.inference.decode src/maxtext/configs/base.yml \
  run_name=decode \
  model_name=deepseek2-16b \
  tokenizer_type=huggingface \
  tokenizer_path=deepseek-ai/DeepSeek-V2-Lite \
  load_parameters_path=gs://ml-auto-solutions/output/unowned/maxtext_nightly_deepseek2-16b-v5p-8-2026-07-01-07-56-03/unscanned/0/items \

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 12.50000% with 7 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/nnx_decoders.py	12.50%	7 Missing ⚠️

📢 Thoughts on this report? Let us know!

[xibinliu]

Replaced with PR#4504, and PR#4480.