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.