Mesa/nnx set defaults true#4253
Draft
mesakhcienet wants to merge 8 commits into
Draft
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).
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
mesakhcienet
force-pushed
the
mesa/nnx-set-defaults-true
branch
from
4e4874e to
abfa824
Compare
…in_step. Previously, test_train_step_cache_hit asserted that the compilation cache contained exactly 1 file, assuming only would be cached. However, in some environments (like nightly workflows), JAX also caches other compiled functions (like ), causing the test to fail. This fix changes the assertion to filter the cache files for the prefix before verifying that only one compilation occurred, making the test robust against other functions being cached.
… Python scalar` that occurs in `RemoteIteratorWrapper` during state save/restore on multi-device topologies (size > 1). It also adds validation to ensure colocated Python data input is only used with Pathways (single controller) enabled, and replaces incorrect usages of `jax.local_devices()` with `global_mesh.devices`. 1. **ValueError in save/restore**: `RemoteIteratorWrapper.save_state` and `restore_state` were attempting to shape the step value array using `self.dummy_array.shape` and shard it across devices. On topologies with more than 1 device, this resulted in a partitioned array. When this partitioned array was passed to the local iterator, attempting to unpack it to a Python scalar (e.g. via `.item()` or direct conversion) failed because JAX does not allow converting partitioned arrays of size > 1 to Python scalars. 2. **Incorrect Device Resolution**: `RemoteIteratorWrapper` was using `jax.local_devices()` to determine CPU/TPU devices. Under Pathways (single-controller), all devices in the cluster are virtualized as local to the JAX client, meaning `jax.local_devices()` returns all devices (including inactive ones during elastic scale-down), which is incorrect for sharding and shape calculations. 3. **Missing Validation**: `colocated_python_data_input` relies on Pathways single-controller mode, but there was no validation enforcing this constraint, which could lead to cryptic failures if misconfigured. 1. **Replicated Scalar for Step**: Modified `RemoteIteratorWrapper.save_state` and `restore_state` in `multihost_dataloading.py` to pass the training step as a replicated 0D JAX scalar array (global shape `()`) with replicated sharding (`NamedSharding` with `PartitionSpec()`). This ensures the array has size 1 on all devices and can be safely converted to a Python scalar by the local iterator. 2. **Use Global Mesh Devices**: Replaced `jax.local_devices()` with `global_mesh.devices` (via `tuple(global_mesh.devices.flat)`) in `RemoteIteratorWrapper.__init__` to ensure it only uses the active devices defined by the global mesh, handling elastic scaling correctly. 3. **Config Validation**: Added a check in `types.py` to raise a `ValueError` if `colocated_python_data_input` is enabled but `enable_single_controller` is false. Added new unit tests in `third_party/py/maxtext/tests/unit/multihost_dataloading_test.py` to verify the fixes: 1. `test_remote_iterator_wrapper_save_state`: Parameterized over different mesh shapes (1, 2, and 4 devices). Instantiates `RemoteIteratorWrapper` and verifies that calling `save_state` successfully writes the state to a JSON file without raising `ValueError`. 2. `test_remote_iterator_wrapper_restore_state`: Parameterized over different mesh shapes. Verifies that `restore_state` successfully restores the state from a JSON file and resumes iteration correctly. These tests are configured to run with `XLA_FLAGS="--xla_force_host_platform_device_count=4"` via the `BUILD` target to simulate multi-device environments. Before submitting this PR, please make sure (put X in square brackets): - [X] I have performed a self-review of my code. For an optional AI review, add the `gemini-review` label. - [X] I have necessary comments in my code, particularly in hard-to-understand areas. - [X] I have run end-to-end tests tests and provided workload links above if applicable. - [X] I have made or will make corresponding changes to the doc if needed. PiperOrigin-RevId: 936996365
This reverts commit 2c062a9.
mesakhcienet
force-pushed
the
mesa/nnx-set-defaults-true
branch
from
abfa824 to
d4061d5
Compare
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
5 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.