Skip to content

Commit 8c1ead8

Browse files
jeffcarpjeffcarp
committed
Add OSS docs, pyproject.toml, and update README
PiperOrigin-RevId: 862885920
1 parent b57cbbd commit 8c1ead8

File tree

4 files changed

+243
-5
lines changed

4 files changed

+243
-5
lines changed

‎README.md‎

Lines changed: 70 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,73 @@
11
# Parallax
22

3-
Parallax is a library for automatically scaling JAX models. You can choose from
4-
various types of parallelism without any modifications to the underlying model.
3+
**Parallax** is a library for automatically scaling JAX models. It simplifies
4+
the process of training large models by offering automatic parallelism
5+
strategies and memory optimizations, allowing you to focus on your model
6+
architecture rather than sharding configurations.
57

6-
This is not an officially supported Google product. This project is not
7-
eligible for the [Google Open Source Software Vulnerability Rewards
8-
Program](https://bughunters.google.com/open-source-security).
8+
**Parallax helps you:**
9+
10+
* **Automatically shard** your JAX models and functions without manually
11+
defining `PartitionSpec`s.
12+
* **Apply advanced parallelism** strategies like Fully Sharded Data Parallel
13+
(FSDP) and Distributed Data Parallel (DDP) with ease.
14+
* **Optimize memory usage** through model offloading (keeping weights in CPU
15+
RAM) and rematerialization.
16+
17+
With Parallax, you can scale off-the-shelf JAX models to run on larger hardware
18+
configurations or fit larger models on existing hardware without extensive code
19+
modifications.
20+
21+
> This is not an officially supported Google product. This project is not
22+
> eligible for the
23+
> [Google Open Source Software Vulnerability Rewards Program](
24+
https://bughunters.google.com/open-source-security).
25+
26+
## Installation
27+
28+
You can install Parallax using pip:
29+
30+
```bash
31+
pip install google-parallax
32+
```
33+
34+
## Usage
35+
36+
Parallax integrates seamlessly with [Flax
37+
NNX](https://flax.readthedocs.io/en/latest/index.html) models. Here is a simple
38+
example of how to use auto-sharding:
39+
40+
```python
41+
import parallax
42+
from flax import nnx
43+
import jax
44+
import jax.numpy as jnp
45+
46+
model = parallax.create_sharded_model(
47+
model_or_fn=lambda: Model(...),
48+
sample_inputs=(jnp.ones((1, 32)),),
49+
strategy=parallax.ShardingStrategy.AUTO,
50+
data_axis_name='fsdp',
51+
model_axis_name='tp',
52+
)
53+
```
54+
55+
## Features
56+
57+
* **AutoSharding**: Automatically discover optimal sharding strategies.
58+
* **FSDP & DDP**: Ready-to-use implementations of common parallel training
59+
strategies.
60+
* **Model Offloading**: Stream model weights from CPU to device memory to
61+
train larger models.
62+
* **Rematerialization**: Automatic activation recomputation to save memory.
63+
64+
## Contributing
65+
66+
We welcome contributions! Please check
67+
[`docs/contributing.md`](docs/contributing.md) for details on how to submit
68+
pull requests and report bugs.
69+
70+
## Support
71+
72+
If you encounter any issues, please report them on our [GitHub
73+
Issues](https://github.com/google/parallax/issues) page.

‎docs/code-of-conduct.md‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
# Code of Conduct
2+
3+
## Our Pledge
4+
5+
In the interest of fostering an open and welcoming environment, we as
6+
contributors and maintainers pledge to making participation in our project and
7+
our community a harassment-free experience for everyone, regardless of age, body
8+
size, disability, ethnicity, gender identity and expression, level of
9+
experience, education, socio-economic status, nationality, personal appearance,
10+
race, religion, or sexual identity and orientation.
11+
12+
## Our Standards
13+
14+
Examples of behavior that contributes to creating a positive environment
15+
include:
16+
17+
* Using welcoming and inclusive language
18+
* Being respectful of differing viewpoints and experiences
19+
* Gracefully accepting constructive criticism
20+
* Focusing on what is best for the community
21+
* Showing empathy towards other community members
22+
23+
Examples of unacceptable behavior by participants include:
24+
25+
* The use of sexualized language or imagery and unwelcome sexual attention or
26+
advances
27+
* Trolling, insulting/derogatory comments, and personal or political attacks
28+
* Public or private harassment
29+
* Publishing others' private information, such as a physical or electronic
30+
address, without explicit permission
31+
* Disrespecting the community's time by sending spam or other unsolicited
32+
commercial messages
33+
* Other conduct which could reasonably be considered inappropriate in a
34+
professional setting
35+
36+
## Our Responsibilities
37+
38+
Project maintainers are responsible for clarifying the standards of acceptable
39+
behavior and are expected to take appropriate and fair corrective action in
40+
response to any instances of unacceptable behavior.
41+
42+
Project maintainers have the right and responsibility to remove, edit, or reject
43+
comments, commits, code, wiki edits, issues, and other contributions that are
44+
not aligned to this Code of Conduct, or to ban temporarily or permanently any
45+
contributor for other behaviors that they deem inappropriate, threatening,
46+
offensive, or harmful.
47+
48+
## Scope
49+
50+
This Code of Conduct applies both within project spaces and in public spaces
51+
when an individual is representing the project or its community. Examples of
52+
representing a project or community include using an official project e-mail
53+
address, posting via an official social media account, or acting as an appointed
54+
representative at an online or offline event. Representation of a project may be
55+
further defined and clarified by project maintainers.
56+
57+
This Code of Conduct also applies outside the project spaces when the Project
58+
Steward has a reasonable belief that an individual's behavior may have a
59+
negative impact on the project or its community.
60+
61+
## Conflict Resolution
62+
63+
We do not believe that all conflict is bad; healthy debate and disagreement
64+
often yield positive results. However, it is never okay to be disrespectful or
65+
to engage in behavior that violates the project’s code of conduct.
66+
67+
If you see someone violating the code of conduct, you are encouraged to address
68+
the behavior directly with those involved. Many issues can be resolved quickly
69+
and easily, and this gives people more control over the outcome of their
70+
dispute. If you are unable to resolve the matter for any reason, or if the
71+
behavior is threatening or harassing, report it. We are dedicated to providing
72+
an environment where participants feel welcome and safe.
73+
74+
Reports should be directed to parallax-jax@google.com, the
75+
Project Steward(s) for Parallax. It is the Project Steward’s duty to
76+
receive and address reported violations of the code of conduct. They will then
77+
work with a committee consisting of representatives from the Open Source
78+
Programs Office and the Google Open Source Strategy team. If for any reason you
79+
are uncomfortable reaching out to the Project Steward, please email
80+
opensource@google.com.
81+
82+
We will investigate every complaint, but you may not receive a direct response.
83+
We will use our discretion in determining when and how to follow up on reported
84+
incidents, which may range from not taking action to permanent expulsion from
85+
the project and project-sponsored spaces. We will notify the accused of the
86+
report and provide them an opportunity to discuss it before any action is taken.
87+
The identity of the reporter will be omitted from the details of the report
88+
supplied to the accused. In potentially harmful situations, such as ongoing
89+
harassment or threats to anyone's safety, we may take action without notice.
90+
91+
## Attribution
92+
93+
This Code of Conduct is adapted from the Contributor Covenant, version 1.4,
94+
available at
95+
https://www.contributor-covenant.org/version/1/4/code-of-conduct/

‎docs/contributing.md‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# How to contribute
2+
3+
We'd love to accept your patches and contributions to this project.
4+
5+
## Before you begin
6+
7+
### Sign our Contributor License Agreement
8+
9+
Contributions to this project must be accompanied by a
10+
[Contributor License Agreement](https://cla.developers.google.com/about) (CLA).
11+
You (or your employer) retain the copyright to your contribution; this simply
12+
gives us permission to use and redistribute your contributions as part of the
13+
project.
14+
15+
If you or your current employer have already signed the Google CLA (even if it
16+
was for a different project), you probably don't need to do it again.
17+
18+
Visit <https://cla.developers.google.com/> to see your current agreements or to
19+
sign a new one.
20+
21+
### Review our community guidelines
22+
23+
This project follows
24+
[Google's Open Source Community Guidelines](https://opensource.google/conduct/).
25+
26+
## Contribution process
27+
28+
### Code reviews
29+
30+
All submissions, including submissions by project members, require review. We
31+
use GitHub pull requests for this purpose. Consult
32+
[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
33+
information on using pull requests.

‎pyproject.toml‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
[build-system]
2+
requires = ["hatchling"]
3+
build-backend = "hatchling.build"
4+
5+
[project]
6+
name = "google-parallax"
7+
version = "0.0.1"
8+
authors = [
9+
{ name="Jeff Carpenter", email="code@jeffcarp.com" },
10+
{ name="Stella Yan", email="sueyan@gmail.com" },
11+
]
12+
description = "A library for autoscaling JAX models."
13+
readme = "README.md"
14+
requires-python = ">=3.8"
15+
classifiers = [
16+
"Programming Language :: Python :: 3",
17+
"Operating System :: OS Independent",
18+
]
19+
license = {file = "LICENSE"}
20+
dependencies = [
21+
"flax",
22+
"jax",
23+
]
24+
25+
[tool.hatch.build]
26+
include = [
27+
"src/*"
28+
]
29+
exclude = [
30+
"**/*_test.py"
31+
]
32+
33+
[tool.hatch.build.targets.wheel]
34+
packages = ["parallax"]
35+
36+
[project.urls]
37+
Homepage = "https://github.com/google/parallax"
38+
Issues = "https://github.com/google/parallax/issues"
39+
40+
[project.optional-dependencies]
41+
dev = [
42+
"absl-py",
43+
"numpy",
44+
"optax",
45+
]

0 commit comments

Comments
 (0)