Microsoft MX Format Emulation Tool on Ubuntu 22.04
Published:
Microscaling (MX) data formats are designed to reduce computational and storage costs in deep learning applications by combining a per-block scaling factor with narrow floating-point and integer types for individual elements. This approach balances hardware efficiency, model accuracy, and user convenience. Empirical results from popular benchmarks demonstrate the practicality of MX data formats as a drop-in replacement for standard FP32 in AI inference and training, with minimal user friction. Notably, MX formats have enabled the training of generative language models with sub-8-bit weights, activations, and gradients, incurring only minor accuracy losses without requiring modifications to the training process.
In 2023, the Open Compute Project (OCP) released the Microscaling Formats (MX) Specification, representing an initial effort to standardize an open, interoperable family of data formats with a shared, fine-grained block scale. This standardization effort includes contributions from industry leaders such as AMD, Arm, Intel, Meta, Microsoft, NVIDIA, and Qualcomm, aiming to create a unified approach to narrow precision data formats in AI workloads. To facilitate the adoption of MX data formats, Microsoft has developed a PyTorch emulation library that provides drop-in replacements for standard PyTorch modules and functions.
In this blog, I will go throught the step-by-step process of installing the Microxcaling repo using docker images. (repo link: https://github.com/microsoft/microxcaling)
Below are the desktop configurations.
OS: Ubuntu 20.04
NVIDIA-SMI 550.54.14
Driver Version: 550.54.14
CUDA Version: 12.4
Steps
- Install docker
$sudo apt update $sudo apt install -y docker.io
- Install the docker image
$sudo docker pull nvcr.io/nvidia/pytorch:24.06-py3
- Install Install NVIDIA Container Toolkit (if not already installed)
$distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \ && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - \ && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list $sudo apt update $sudo apt install -y nvidia-container-toolkit
- Log in the docker image
$sudo docker run --gpus all -it --rm -v $(pwd):/workspace nvcr.io/nvidia/pytorch:24.06-py3
- Update pip
$python -m pip install --upgrade pip
- Clone the repo. Update the requirements.txt file with the following versions.
torch==2.4.1 torchvision==0.19.1 torchaudio==2.4.1
- Install required packages
$pip install -r requirements.txt
- Set the PYTHONPATH Manually. If no errors occur, the path is set correctly.
$export PYTHONPATH=/workspace/microxcaling:$PYTHONPATH
$python -c "import mx"
- Run an exmaple test
/workspace/microxcaling/examples# sh run_mxfp6.sh