NNX: fix Linen-parity gaps on the default path + unit tests#4255
Draft
ecnal-cienet wants to merge 4 commits into
Draft
NNX: fix Linen-parity gaps on the default path + unit tests#4255ecnal-cienet wants to merge 4 commits into
ecnal-cienet wants to merge 4 commits into
Conversation
PR6-PR10 promoted every routed-to-Linen feature to NNX-native; PR#2885 lands NNX-native pipeline parallelism. This PR flips the three defaults in base.yml so NNX is the production path, and bundles the NNX-only fixes that surface once pure_nnx=True (DiLoCo merge/checkpoint, Zero-1 input shardings on flat nnx.State, MTP sown-Variable handling, generate_param_only_checkpoint NNX flow, maxengine Linen-parity removal).
ecnal-cienet
force-pushed
the
fix/nnx-linen-parity-gaps
branch
from
0390217 to
aa18ab3
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Features that were Linen-only or silently divergent on the now-default NNX path: - skip_step_on_spikes: forward loss/grad_norm through the NNX optimizer; surface step_skipped - loss_fn: match Linen indexer-warmup vs vocab-tiling ordering - decoder logits guards: use the call-arg model_mode, not self.model_mode - routed_bias: read moe_bias_updates by suffix (NNX intermediates are model-rooted) - record_internal_nn_metrics: NNX-aware activation-metric collection (was KeyError) - qwix: skip non-Variable attrs in the ToNNX bridge; guard pure_nnx + bridged decoder + qwix quant - maxengine.set_engine_vars_from_base_engine: dispatch to NNX kv-cache annotations - GRPO (NNX): gradient_accumulation_steps>1, scan_layers=False reshard, optimizer_memory_host_offload
ecnal-cienet
force-pushed
the
fix/nnx-linen-parity-gaps
branch
from
aa18ab3 to
bdf087f
Compare
grpo_trainer: the GA loss metric divided sum(total_loss) by total_weights and re-added moe_lb, but total_loss is already a per-batch mean that includes moe_lb; report sum(total_loss)/GA so it matches the full-batch step (gradients were already correct). Tests: - GRPO grad-accum equivalence: GA=2 reproduces the full-batch params and loss - skip_step_on_spikes: a real spike is skipped and params frozen - routed_bias: moe_bias_updates found by suffix on the model-rooted NNX dict - record_activation_metrics: identical metrics for Linen- and NNX-shaped inputs - qwix quant guard: rejects bridged-decoder+qwix, accepts the pure-decoder path
…ffload (#10) Both paths need Pathways/TPU-memory infra at runtime, so the external pieces (reshard_pytree via pathwaysutils; move_memory_to_device) are mocked and the test pins our changes: - #9: scan_layers=False no longer raises and the unscanned policy params are pushed to the inference engine (guard removal). - #10: optimizer_memory_host_offload runs the device_put/update plumbing and yields the same params as the no-offload step (memory placement, not math).
ecnal-cienet
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This closes the Linen-only / silently-divergent gaps on that path and adds correctness unit tests.
Stacked on
feat/nnx-set-defaults-true.Fixes
skip_step_on_spikes: silent no-op on NNX (apply_gradientsdidn't forwardloss/grad_norm). Now forwarded + metric surfaced.loss_fn: NNX checked vocab-tiling before indexer warm-up; reordered to match Linen.self.model_modeinstead of the call-argmodel_mode.routed_bias: updates silently dropped on NNX (Linen"intermediates"prefix absent on NNX dict). Now matched by suffix.record_internal_nn_metrics: KeyError on NNX. Now NNX-aware via suffix collection.pure_nnxwith the bridged decoder; bridge now skips qwix's non-Variableattrs + a config guard rejects bridged-decoder+qwix.maxengine.set_engine_vars_from_base_engine: AttributeError on NNX; now usesget_kv_cache_annotations_nnx.gradient_accumulation_steps>1: NotImplementedError on NNX. Implemented; also fixed the GA loss metric (sum/GA, notsum/total_weights).scan_layers=False: NotImplementedError on NNX. Guard removed (NNX policy already matches the inference layout).optimizer_memory_host_offload: ignored on NNX; now moves opt state to device before the update.Also re-declared the legacy GRPO config fields (
inference_replicas/inference_devices_per_replica/inference_rollouts/use_pathways_reshard) intypes.py— they were dropped from the schema sogrpo.ymlcouldn't load (pre-existing onmain).Tests
tests/unit/{train_nnx_test,grpo_nnx_test,maxengine_nnx_test,nnx_quant_guard_test}.py— 27 pass on CPU:Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.