fix(nnx): keep only newly created params in _apply_layers_sequentially

[hodaaaaaaaaaa]

Description

Under dynamic_graph_init, the scan body in NNXDecoder._apply_layers_sequentially returns the full nnx.Param state. jax.lax.scan stacks every output, so each Qwix LoRA injection materializes a fresh copy of the whole layer stack: on gemma2-2b, 1194 MB per adapter, for adapter weights of 6.1 MB. Nothing errors, but the base is no longer shared between adapters, so multi-adapter serving on one worker is impossible with scan on.

The branch is necessary — new LoRA params are created inside the scan body and this is how they get out. It just carries the base out with them.

5e3a353 ("fix(nnx): keep params created inside the layer scan") fixed the same problem in nnx_scan.apply_scanned_layers by returning only params created inside the body. That applier has two callers (Gemma 4, Qwen3-Next); everything else scanned, gemma2 included, goes through _apply_layers_sequentially, which never got the fix because it doesn't crash.

This PR transcribes that fix: record the paths fed in as scan inputs, return only paths not in that set, and read the base back off the pre-scan layers module with nnx.state(layers, nnx.Param). That call preserves array identity, so the base is genuinely shared rather than just smaller, and it returns at its original param_scan_axis layout.

One semantic change: under dynamic_graph_init, in-body mutations of carried params are no longer returned. The non-dynamic branch already discards them and apply_scanned_layers assumes the same — the applier already treats scanned params as read-only inside the body.

The two appliers are slated to be unified (see the note at nnx_decoders.py:1024); fixing this first means the merge won't have to reconcile two different behaviors.

FIXES: #5129

Tests

New regression test — TestApplyLayersSequentiallyDynamicGraphInit in tests/unit/nnx_decoders_test.py, following the existing TestApplyLayersSequentiallyMetadataAxisName harness. It drives _apply_layers_sequentially directly with an nnx.vmap-stacked dummy layer that creates an nnx.LoRAParam while tracing, and asserts that the new param escapes the scan and that the base array is the same object afterwards, at param_scan_axis 0 and 1. The identity assertion fails on main at both axes and passes with this change.

Existing suites, on CPU:

XLA_FLAGS=--xla_force_host_platform_device_count=8 JAX_PLATFORMS=cpu \
  pytest tests/unit/nnx_decoders_test.py tests/unit/nnx_decoder_test.py \
         tests/unit/nnx_scan_test.py tests/unit/lora_utils_nnx_test.py \
         tests/unit/nnx_wrappers_test.py
# 75 passed, 2 skipped

End-to-end check — build gemma2-2b with base_num_decoder_layers=2, scan_layers=True, an fsdp=4 × tensor=2 mesh over 8 forced host devices, then call lora_utils.apply_lora_to_model(model, mesh, config) twice and diff jax.live_arrays() around each call:

	before	after
per-adapter memory	1194.2 MB	6.1 MB
base array identity	2/24 same object	24/24
LoRA factors created	28	28
nnx.jit forward	passes	passes

Adapter outputs are bitwise identical before and after (np.array_equal, compared across processes): the change affects memory, not numerics. The two arrays already shared before the fix are token_embedder/embedding and decoder/decoder_norm/scale, the only Gemma 2 weights outside the layer stack.

All of the above runs on plain CPU; no TPU required.

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. — not run; verified with the unit suites and the CPU multi-adapter check above. Happy to run more if a maintainer points me at the right target.
• 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. — no doc changes needed.

[gemini-code-assist]

Code Review

This pull request modifies _apply_layers_sequentially in nnx_decoders.py to prevent carried-in parameters from being re-emitted as scan outputs during dynamic graph initialization, while ensuring parameters created inside the scan body (such as LoRA adapters) are correctly returned. It also adds unit tests to verify this behavior. The reviewer identified a critical issue where this change exposes a bug in _apply_interleaved_scanned_layers, as the returned chunk_stack is not written back to self, which would cause newly created parameters inside interleaved layers to be lost.