MPX: Mixed Precision Training for JAX

TL;DR

MPX is a new toolbox that enables efficient mixed-precision training in JAX, simplifying integration with existing frameworks and maintaining model accuracy through dynamic loss-scaling and type management.

Contribution

It introduces MPX, the first comprehensive mixed-precision training toolkit for JAX, with seamless integration, dynamic loss-scaling, and precision management features.

Findings

Enables mixed-precision training in JAX with minimal code changes.

Maintains model accuracy through dynamic loss-scaling.

Provides automatic gradient and optimizer management for mixed precision.

Abstract

Mixed-precision training has emerged as an indispensable tool for enhancing the efficiency of neural network training in recent years. Concurrently, JAX has grown in popularity as a versatile machine learning toolbox. However, it currently lacks robust support for mixed-precision training. We propose MPX, a mixed-precision training toolbox for JAX that simplifies and accelerates the training of large-scale neural networks while preserving model accuracy. MPX seamlessly integrates with popular toolboxes such as Equinox and Flax, allowing users to convert full-precision pipelines to mixed-precision versions with minimal modifications. By casting both inputs and outputs to half precision, and introducing a dynamic loss-scaling mechanism, MPX alleviates issues like gradient underflow and overflow that commonly arise in half precision computations. Its design inherits critical features from JAX's type-promotion behavior, ensuring that operations take place in the correct precision and allowing for selective enforcement of full precision where needed (e.g., sums, means, or softmax). MPX further provides wrappers for automatic creation and management of mixed-precision gradients and optimizers, enabling straightforward integration into existing JAX training pipelines. MPX's source code, documentation, and usage examples are available at github.com/Data-Science-in-Mechanical-Engineering/mixed_precision_for_JAX .