Adding AxoNN's 3D tensor parallelism [WIP]

[siddharth9820]

Steps to run -

Install AxoNN (dependencies - Pytorch and mpi4py) -

• git clone [email protected]:axonn-ai/axonn.git
• cd axonn
• git checkout 45647ea
• pip install -e .

Preparing a config file to use AxoNN -

First, set "use_axonn_model_parallelism": true
Then set "depth_model_parallel_size", "row_model_parallel_size", "column_model_parallel_size" as per the requirements of your model. The product of these should equal "model_parallel_size".

you can also set "optimize_axonn_communication: true" to enable communication optimizations. These also require you to set the environment variable - export CUDA_DEVICE_MAX_CONNECTIONS=1

At a high level, the matrix multiplications in your model will be sharded over "model_parallel_size" GPUs.

ToDos -

• Implement axonn support for parallel_output=True (gpt j residual)
• Implement axonn support for LLaMAMLP
• Integrate communication optimizations

[Quentin-Anthony]

Under testing

[siddharth9820]

Correctness check on 125M.yml with
use_axonn_model_parallelism:true, column_model_parallel_size=1, row_model_parallel_size=1, depth_model_parallel_size=2, model_parallel_size=2 on 2 GPUs.
Dataset - enwiki8

(the loss curve is smoothed over 100 iterations)

[siddharth9820]

@Quentin-Anthony I have updated the install instructions to install axonn from a fixed commit - 3ebc34c

[Quentin-Anthony]

@Quentin-Anthony I have updated the install instructions to install axonn from a fixed commit - 3ebc34c

Thanks!

[siddharth9820]

@Quentin-Anthony Pushed some communication optimizations and also updated the instructions to install axonn from a newer commit - 45647ea.

To enable these optimizations, you just need to set "optimize_axonn_communication: true" in your neo-x config files. These also require you to set the environment variable - export CUDA_DEVICE_MAX_CONNECTIONS=1