Developer Guide

Code Organization

The code in alpa’s repository is organized as follows:

alpa: the python source code of Alpa
benchmark: benchmark scripts
build_jaxlib: build scripts for Alpa’s version of jaxlib
docs: documentation and tutorials
examples: public examples
playground: experimental scripts
tests: unit tests
third_party: third party repos

In addition, Alpa maintains a tensorflow fork. This is because alpa modifies the XLA compiler, whose code is hosted in the tensorflow repo.

tensorflow-alpa: The TensorFlow fork for Alpa. The c++ source code of Alpa mainly resides in tensorflow/compiler/xla/service/spmd.

Contribute to Alpa

Please submit a pull request if you plan to contribute to Alpa.

Formatting and Linting

We follow Google Python Style Guide.

Install yapf and pylint via:

pip install yapf==0.32.0 pylint==2.14.0

Use the following script to format the code and check linting errors:

./format.sh

Unit Testing

Every New feature should come with a unit test. See this README.md on how to run tests locally.

Updating submodule tensorflow-alpa

Alpa repo stores a commit hash of the submodule tensorflow-alpa, so git knows which version of tensorflow-alpa should be used. However, commands like git pull do not update the submodule to the latest stored commit. You need to additionally use the commands below.

git submodule update --init --recursive

Contributing to submodule tensorflow-alpa

If you want to contribute code to tensorflow-alpa, you can follow the steps below

Contributors send a pull request to tensorflow-alpa.
Maintainers review the pull request and merge it to tensorflow-alpa.
Contributors send a pull request to alpa. The pull request should update the stored hash commit of the submodule and other modifications to alpa if necessary.
Maintainers review the pull request and merge it to alpa.