Alpa automatically parallelizes large tensor computation graphs and runs them on a distributed cluster.
Project description
Alpa
Alpa is a system for large-scale distributed training. Alpa is specifically designed for training giant neural networks that cannot fit into a single device. Alpa can automatically generate dstirbuted execution plans that unify data, operator, and pipeline parallelism.
Quick Start
Use Alpa's decorator @parallelize
to scale your single-node training code to distributed clusters, even though
your model is much bigger than a single device memory.
import alpa
@alpa.parallelize
def train_step(model_state, batch):
def loss_func(params):
out = model_state.forward(params, batch["x"])
return jnp.mean((out - batch["y"]) ** 2)
grads = grad(loss_func)(model_state.params)
new_model_state = model_state.apply_gradient(grads)
return new_model_state
# The training loop now automatically runs on your designated cluster.
model_state = create_train_state()
for batch in data_loader:
model_state = train_step(model_state, batch)
Check out the Alpa Documentation site for installation instructions, tutorials, examples, and more.
More Information
- Alpa paper (OSDI'22)
- Blog
Contributing
Please read the contributor guide if you are interested in contributing to Alpa. Please connect to Alpa contributors via the Alpa slack.
License
Alpa is licensed under the Apache-2.0 license.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distributions
Hashes for alpa-0.0.1-cp39-cp39-manylinux2014_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 68f183edd100f44a03272cd7732c9dbf3e52d1a7a86c09a86190858ee1913e4e |
|
MD5 | d6e07c8894d63405445186b7597cd1e1 |
|
BLAKE2b-256 | 766fdda19b2ff125d91a4cce60018a1f9c1d740171c65ad5e91de767a201abb5 |
Hashes for alpa-0.0.1-cp38-cp38-manylinux2014_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | a46b8a7ba79944a028e8caeae5febf6ab201503fb2008dc5f3a52388911ef008 |
|
MD5 | dffaf27097494f8e2821ea34bdf05599 |
|
BLAKE2b-256 | 9b3a1dd46eaabd73d069c4729c6fa1f08eda361ff81a08f04776c1c9b5078004 |
Hashes for alpa-0.0.1-cp37-cp37m-manylinux2014_x86_64.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 6ec751e1687dd2a9e4eb959dd10db47761fbb3a2247f9bd403d76f394c978598 |
|
MD5 | 8208499ea9f78473a766cb877a00153a |
|
BLAKE2b-256 | 21d364565b0af48f133ef81fafcabccce9a21c5fa51273b45247e231abc825eb |