Oscillation-Reduced MXFP4 Training for Vision Transformers

TL;DR

This paper introduces TetraJet, a novel training method for FP4 precision in Vision Transformers, which reduces accuracy loss caused by weight oscillation and achieves performance close to full precision training.

Contribution

The paper proposes TetraJet with Q-EMA and Q-Ramping to address weight oscillation in MXFP4 training, significantly improving accuracy over existing methods.

Findings

TetraJet reduces accuracy degradation by over 50%.

Q-EMA and Q-Ramping effectively mitigate oscillation.

Achieves competitive performance with full precision training.

Abstract

Pre-training Transformers in FP4 precision is becoming a promising approach to gain substantial speedup, but it comes with a considerable loss of accuracy. Microscaling (MX) data format provides a fine-grained per-group quantization method to improve the representation ability of the FP4 format and is supported by the next-generation Blackwell GPU architecture. However, training with MXFP4 data format still results in significant degradation and there is a lack of systematic research on the reason. In this work, we propose a novel training method TetraJet for a more accurate FP4 training. We comprehensively evaluate all of the quantizers involved in the training, and identify the weight oscillation problem in the forward pass as the main source of the degradation in MXFP4 training. Therefore, we introduce two novel methods, EMA Quantizer (Q-EMA) and Adaptive Ramping Optimizer (Q-Ramping), to resolve the oscillation problem. Extensive experiments on Vision Transformers demonstrate that TetraJet consistently outperforms the existing 4-bit training methods, and Q-EMA & Q-Ramping can provide additional enhancement by effectively reducing oscillation. We decreased the accuracy degradation by more than $50\%$ compared to the baseline, and can even achieve competitive performance compared to full precision training. The codes are available at https://github.com/thu-ml/TetraJet-MXFP4Training