From 2995062492e7b74d6e8d6b2b6a20dca487601c2a Mon Sep 17 00:00:00 2001 From: Hana Joo Date: Wed, 1 Jul 2026 07:30:28 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 941083877 --- .github/ISSUE_TEMPLATE/bug_report.md | 22 +++ .github/PULL_REQUEST_TEMPLATE.md | 6 + .gitignore | 216 +++++++++++++++++++++++++++ parallax/offload.py | 10 +- parallax/offload_test.py | 4 +- parallax/sharding/auto_shard.py | 18 ++- parallax/sharding/auto_shard_test.py | 43 +++++- parallax/sharding/base.py | 12 +- parallax/sharding/base_test.py | 15 +- 9 files changed, 325 insertions(+), 21 deletions(-) create mode 100644 .github/ISSUE_TEMPLATE/bug_report.md create mode 100644 .github/PULL_REQUEST_TEMPLATE.md create mode 100644 .gitignore diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000..fee303a --- /dev/null +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,22 @@ +--- +name: Bug report +about: Create a report to help us improve + +--- + +## Expected Behavior + + +## Actual Behavior + + +## Steps to Reproduce the Problem + +1. +1. +1. + +## Specifications + +- Version: +- Platform: \ No newline at end of file diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md new file mode 100644 index 0000000..ba31ec0 --- /dev/null +++ b/.github/PULL_REQUEST_TEMPLATE.md @@ -0,0 +1,6 @@ +Fixes # + +> It's a good idea to open an issue first for discussion. + +- [ ] Tests pass +- [ ] Appropriate changes to documentation are included in the PR diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e15106e --- /dev/null +++ b/.gitignore @@ -0,0 +1,216 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +# Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry.lock +# poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +# pdm.lock +# pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +# pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# Redis +*.rdb +*.aof +*.pid + +# RabbitMQ +mnesia/ +rabbitmq/ +rabbitmq-data/ + +# ActiveMQ +activemq-data/ + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +# .idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml diff --git a/parallax/offload.py b/parallax/offload.py index e0b14f3..0752510 100644 --- a/parallax/offload.py +++ b/parallax/offload.py @@ -104,7 +104,7 @@ def forward(model_chunk, intermediate_result): # example element 0 is the inputs to layer 0. The last element of the array # are the final outputs (which become the inputs to the loss function). saved_intermediates = [inputs] - for chunk in model.layers: + for chunk in model.layers: # pyrefly: ignore[missing-attribute] logits = forward(chunk, saved_intermediates[-1]) saved_intermediates.append(logits) @@ -132,14 +132,14 @@ def offload_backward( A full nnx.State mapping representing a composite of the grads from all layers. """ - if len(saved_intermediates) != len(model.layers) + 1: + if len(saved_intermediates) != len(model.layers) + 1: # pyrefly: ignore[missing-attribute] raise ValueError( 'The length of `saved_intermediates` must match the number of model ' 'layers plus one.' ) layer_grads = {} - for i, chunk in reversed(list(enumerate(model.layers))): + for i, chunk in reversed(list(enumerate(model.layers))): # pyrefly: ignore[missing-attribute] # pylint: disable=cell-var-from-loop chunk_input = saved_intermediates[i] graphdef, state = nnx.split(chunk) @@ -189,13 +189,13 @@ def remat_model(model: nnx.Module) -> nnx.Module: """Takes an NNX Module and returns one with all layers rematerialized.""" # TODO(jeffcarp): Generalize this to work with non-Sequential models. new_model = nnx.clone(model) - for i, layer in enumerate(new_model.layers): + for i, layer in enumerate(new_model.layers): # pyrefly: ignore[missing-attribute] signature = inspect.signature(layer.__call__) unbound_call = layer.__class__.__call__ # Shift static_argnums by 1 because of `self` in unbound call. static_argnums = tuple(i + 1 for i in range(1, len(signature.parameters))) rematted_call = nnx.remat(unbound_call, static_argnums=static_argnums) - new_model.layers[i] = RemattedLayer(layer, rematted_call) + new_model.layers[i] = RemattedLayer(layer, rematted_call) # pyrefly: ignore[missing-attribute] return new_model diff --git a/parallax/offload_test.py b/parallax/offload_test.py index fb22f32..aabd0f8 100644 --- a/parallax/offload_test.py +++ b/parallax/offload_test.py @@ -100,8 +100,8 @@ def loss_fn(model, inputs): np.testing.assert_array_equal(act_loss, exp_loss) # Assert all model weights match after gradient update. np.testing.assert_allclose( - actual_model.layers[0].kernel[...], - reference_model.layers[0].kernel[...], + actual_model.layers[0].kernel[...], # pyrefly: ignore[missing-attribute] + reference_model.layers[0].kernel[...], # pyrefly: ignore[missing-attribute] atol=1e-3, rtol=1e-1, ) diff --git a/parallax/sharding/auto_shard.py b/parallax/sharding/auto_shard.py index 523fcde..67bae60 100644 --- a/parallax/sharding/auto_shard.py +++ b/parallax/sharding/auto_shard.py @@ -86,14 +86,20 @@ def get_shardings( params_assignments = [] for var in jaxpr.jaxpr.invars[: len(params_flat)]: params_assignments.append([]) - for i in range(var.aval.ndim): # pytype: disable=attribute-error + dim_sharded = False + for i in range(var.aval.ndim - 1, -1, -1): # pytype: disable=attribute-error root = graph.get_root((var, i)) - if (model_axis is not None and (root, model_axis) in edges) or var.aval.shape[i] < min_shard_size: # pytype: disable=attribute-error + if ( + (model_axis is None) # pyrefly: ignore[unbound-name] + or ((root, model_axis) in edges) + or (var.aval.shape[i] < min_shard_size) + or dim_sharded + ): params_assignments[-1].append(None) # conflict with model axis - else: - params_assignments[-1].append( - model_axis_name if model_axis is not None else None - ) + else: # Assign model axis if not already sharded + params_assignments[-1].append(model_axis_name) + dim_sharded = True + params_assignments[-1].reverse() params_assignments = params_treedef.unflatten(params_assignments) inputs_assignments = [] diff --git a/parallax/sharding/auto_shard_test.py b/parallax/sharding/auto_shard_test.py index aba885c..4e30b57 100644 --- a/parallax/sharding/auto_shard_test.py +++ b/parallax/sharding/auto_shard_test.py @@ -13,19 +13,33 @@ # limitations under the License. import os +from absl.testing import absltest +from absl.testing import parameterized from flax import nnx import jax import jax.numpy as jnp from parallax.examples import models from parallax.sharding import auto_shard -from absl.testing import absltest -from absl.testing import parameterized - NamedSharding = jax.sharding.NamedSharding P = jax.sharding.PartitionSpec +class ResidualBlockModel(nnx.Module): + + def __init__(self, rngs: nnx.Rngs): + self.linear1 = nnx.Linear(128, 64, rngs=rngs) + self.linear2 = nnx.Linear(64, 32, rngs=rngs) + self.linear3 = nnx.Linear(96, 8, rngs=rngs) + + def __call__(self, x): + z1 = self.linear1(x) + z2 = self.linear2(z1) + z1_z2 = jnp.concatenate([z2, z1], axis=-1) + z3 = self.linear3(z1_z2) + return z3 + + class AutoShardTest(parameterized.TestCase): def setUp(self): @@ -278,6 +292,29 @@ def fn(state, *inputs): self.assertEqual(in_shd, expected_assignments[1]) self.assertEqual(out_shd, expected_assignments[2]) + def test_residual_block_w3_conflict(self): + model = ResidualBlockModel(rngs=nnx.Rngs(0)) + inputs = (jnp.ones((32, 128)),) + graphdef, state = nnx.split(model) + + def fn(state, *inputs): + model = nnx.merge(graphdef, state) + return model(*inputs) + + (model_shd, *in_shd), out_shd = auto_shard.get_shardings( + fn, state, *inputs, min_shard_size=0 + ) + model_shd = nnx.to_pure_dict(model_shd) + + expected_model_shd = { + 'linear1': {'bias': P('model'), 'kernel': P(None, 'model')}, + 'linear2': {'bias': P(None), 'kernel': P('model', None)}, + 'linear3': {'bias': P('model'), 'kernel': P(None, 'model')}, + } + self.assertEqual(model_shd, expected_model_shd) + self.assertEqual(in_shd, [P('data', None)]) + self.assertEqual(out_shd, P('data', None)) + if __name__ == '__main__': # Fake 8 CPUs for testing. diff --git a/parallax/sharding/base.py b/parallax/sharding/base.py index a2fea2a..23cc912 100644 --- a/parallax/sharding/base.py +++ b/parallax/sharding/base.py @@ -37,6 +37,14 @@ class ShardingStrategy(enum.Enum): MANUAL = 'manual' +def _constrain_or_reshard(x: PyTree, sharding: PyTree) -> PyTree: + mesh = jax.sharding.get_abstract_mesh() + if not mesh.empty and mesh.are_all_axes_explicit: + return jax.reshard(x, sharding) + else: + return jax.lax.with_sharding_constraint(x, sharding) + + def create_sharded_model( model_or_fn: nn.Module | Callable[[], nnx.Module], sample_inputs: Any, @@ -77,7 +85,7 @@ def create_sharded_model( model.apply, variables, *sample_inputs, **kwargs ) - sharded_params = jax.lax.with_sharding_constraint(variables, params_shd) + sharded_params = _constrain_or_reshard(variables, params_shd) return sharded_params else: # Support for Flax NNX models. @@ -92,7 +100,7 @@ def fn(state, *inputs): return nnx.merge(graphdef, state)(*inputs) (params_shd, _), _ = get_shardings_fn(fn, state, *sample_inputs, **kwargs) - sharded_state = jax.lax.with_sharding_constraint(state, params_shd) + sharded_state = _constrain_or_reshard(state, params_shd) nnx.update(model, sharded_state) return model diff --git a/parallax/sharding/base_test.py b/parallax/sharding/base_test.py index 91ffac3..2471787 100644 --- a/parallax/sharding/base_test.py +++ b/parallax/sharding/base_test.py @@ -53,7 +53,7 @@ def forward(params, inputs): return model.apply(params, inputs) output = forward(params, x) - np.testing.assert_array_almost_equal(output, reference_output, decimal=6) + np.testing.assert_array_almost_equal(output, reference_output, decimal=6) # pyrefly: ignore[bad-argument-type] # Verify compiled shardings. compiled = forward.lower(params, x).compile() # type: ignore @@ -84,7 +84,7 @@ def forward(params, inputs): return model.apply(params, inputs) output = forward(params, x) - np.testing.assert_array_almost_equal(output, reference_output, decimal=6) + np.testing.assert_array_almost_equal(output, reference_output, decimal=6) # pyrefly: ignore[bad-argument-type] # Verify compiled shardings. compiled = forward.lower(params, x).compile() # type: ignore @@ -115,7 +115,7 @@ def forward(params, inputs): sharded_forward = base.jit(forward, strategy=base.ShardingStrategy.DDP) output = sharded_forward(params, x) - np.testing.assert_array_almost_equal(output, reference_output, decimal=6) + np.testing.assert_array_almost_equal(output, reference_output, decimal=6) # pyrefly: ignore[bad-argument-type] # Verify compiled shardings. compiled = sharded_forward.lower(params, x).compile() # type: ignore @@ -220,6 +220,15 @@ def test_create_sharded_model_nnx(self, strategy, expected_shardings): ) self.assertDictEqual(actual_shardings, expected_shardings) + def test_create_sharded_model_explicit_mesh(self): + mesh = jax.make_mesh((2, 4), ('data', 'model')) + with jax.set_mesh(mesh): + model = base.create_sharded_model( + lambda: models.SimpleMLP(8, 16, 8, rngs=nnx.Rngs(0)), + sample_inputs=(jnp.ones((8, 8)),), + ) + self.assertIsNotNone(model) + def test_sharded_nnx_optimizer_training_loop(self): """Ensures create_sharded_model works within an NNX training loop.""" dummy_inputs, dummy_labels = jnp.ones((4, 16)), jnp.ones((4, 16))