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.
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.
- Alpa paper (OSDI'22)
- Blog
Please read the contributor guide if you are interested in contributing to Alpa. Please connect to Alpa contributors via the Alpa slack.
Alpa is licensed under the Apache-2.0 license.