Fix/nnx linen parity gaps unit test#4267
Draft
hsuan-lun-chiang wants to merge 5 commits into
Draft
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
hsuan-lun-chiang
force-pushed
the
fix/nnx-linen-parity-gaps-unit-test
branch
from
a8db416 to
4c2e1a4
Compare
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
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).
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
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).
hsuan-lun-chiang
force-pushed
the
fix/nnx-linen-parity-gaps-unit-test
branch
from
4c2e1a4 to
d68f45c
Compare
…ing as nnx.Intermediate
This change cleans up transient variables sowed during model creation/tracing from the persistent sharding state and checkpoint layout.
1. Clean Quantization Setup:
- Updated maybe_quantize_model in quantizations.py to pop sowed nnx.Intermediate variables in-place immediately after Qwix tracing.
- This cleans the initial state returned by the model factory fn, allowing us to revert the setup-time nnx.Not(nnx.Intermediate) filters in maxtext_utils.py and train_compile.py.
- Optimized create_nnx_abstract_model in model_creation_utils.py to reuse abs_model.mesh and avoid calling eval_shape twice.
2. Metrics Sowing Uniformity:
- Changed all self.sow("intermediates", ...) calls to self.sow(nnx.Intermediate, ...) in all pure NNX models (gemma.py, llama2.py, etc.).
- Sowing string-based "intermediates" was dynamically creating custom Variable types that did not inherit from nnx.Intermediate, leading to parameter/checkpoint bloat.
- train_step in train.py now strips all intermediates via nnx.Not(nnx.Intermediate) before returning.
3. Testing:
- Added MaybeQuantizeModelTest in quantizations_test.py to assert that sowed intermediates are popped from the model state and that abstract model state contains no intermediates.
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
3 participants
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
Start with a short description of what the PR does and how this is a change from
the past.
The rest of the description includes relevant details and context, examples:
If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456
You can also provide a comma-separated list. If you don't want to close a bug but
simply to reference it, use BUGS, e.g.:
BUGS: b/123456
Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.
Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.
Tests
Please describe how you tested this change, and include any instructions and/or
commands to reproduce.
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.