developer_guide.rst

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

1. Contributors send a pull request to tensorflow-alpa.
2. Maintainers review the pull request and merge it to tensorflow-alpa.
3. 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.
4. Maintainers review the pull request and merge it to alpa.