AxoNN is a parallel framework for training deep neural networks.
Features:
- Offers a hybrid of inter-layer parallelism with pipelining and data parallelism.
- Supports both 16-bit mixed precision and 32-bit full precision training.
- A highly efficient and scalable implementation of inter-layer parallelism with pipelining using asynchronous MPI-based communication and message-driven scheduling that achieves significant overlap of computation and communication.
- Memory optimizations that can reduce the model state memory consumption by 5x for mixed precision training using the Adam optimizer and indirectly also increase hardware efficiency.