Code Structure of the Intra-op Solver
The specific code of the intra-op solver (a.k.a auto-sharding) is scattered in various files of the project. This page contains some pointers to key components of the intra-op solver and help you navigate the complicated code base.
All the links below are based on alpa v0.2.2
- Main entrance:
python entrance (
- Where the possible sharding strategies are registred:
- Where the ILP solver is called:
How to Read and Learn the Code
Run some simple examples
You can run the unit tests under https://github.com/alpa-projects/alpa/tree/v0.2.2/tests/shard_parallel and set break points in the python entrance
You can start from the most basic ones in
Inspect the sharding strategy
You can print the HLO before and after the
How to Debug
Set global environment variable
ALPA_DEBUG_PRINT_AS_STRATEGY=1. This will print the choosen sharding strategy for each instruction and edge costs in a prettier way.