No public description

.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:

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,6 @@
+Fixes #<issue_number_goes_here>
+
+> It's a good idea to open an issue first for discussion.
+
+- [ ] Tests pass
+- [ ] Appropriate changes to documentation are included in the PR

.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

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
 
 

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,
     )

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 = []

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.