anj-s/alpa

Auto parallelization for large-scale neural networks

Alpa

Documentation | Slack

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.