This presentation introduces "Cortet," a research initiative focusing on achieving optimal native FP4 training for large language models (LLMs). It addresses the inherent inefficiency of training LLMs with high-precision floating-point numbers, given that their effective information capacity per parameter is significantly lower. The researchers propose a novel, robust comparison framework built upon scaling laws, which enables accurate prediction of performance and identification of optimal training schemes by normalizing for both inference and training compute costs across various quantization methods and model sizes. Central to their approach are highly optimized custom CUDA kernels for mixed-precision FP4 (MXFP4) operations, including Quest rounding and Hadamard transforms, which minimize computational overhead. The findings demonstrate that native FP4 training, particularly for practical LLM pre-training scenarios, can be optimal, yielding substantial speedups (e.g., ~2.2x for forward pass compared to FP8) while maintaining competitive accuracy, thus outperforming existing 4-bit quantization baselines.
Summarizing the Cortet Research: Optimal Native FP4 Training for LLMs
0:00 Introduction to Cortet: The presentation introduces "Cortet" research, focusing on how native FP4 training can be optimal for Large Language Models (LLMs), featuring Roberto and Andre from ISTA lab.
0:41 The Inefficiency of High Precision: LLMs are typically trained with high-precision floating-point numbers (e.g., FP32), but their actual information capacity is much lower, often around 1-2 bits per parameter, indicating a significant discrepancy.
4:17 Rise of Low Precision Data Types: Hardware manufacturers like Apple, Qualcomm, and Nvidia are increasingly supporting low-precision data types (FP16, BF16, Int8, FP8, FP4) with dedicated kernel support for efficient machine learning operations.
7:07 Quantizing Matrix Multiplications: The work focuses on quantizing matrix multiplication operations within LLMs, which constitute the majority of computations, to explore better accuracy-to-speed trade-offs with lower precision.
8:01 The Accuracy-Speed Trade-off: Reducing precision generally leads to a decrease in model accuracy but offers significant theoretical speedups (e.g., FP4 could be 4x faster than FP16).
12:24 Need for a Robust Comparison Model: Traditional methods often provide limited insights for different model sizes or training durations. A new comparison model is introduced to aggregate data from many training runs, enable extrapolation, and accurately describe the accuracy-speed trade-off.
13:58 Scaling Laws for Quantized Training: A new scaling law is proposed that incorporates forward (P_forward) and backward (P_backward) precisions. P_forward acts as a multiplicative factor on parameter count (affecting model capacity), and P_backward affects the effective number of tokens (influencing convergence speed due to gradient noise).
19:37 De-coupling Inference and Training Costs: The framework distinguishes between inference cost (primarily affected by forward pass quantization) and training cost (primarily affected by backward pass quantization), normalizing these costs to compare different quantization schemes.
23:06 Compute Cost-Normalized Comparisons: By normalizing inference and training costs, the model can predict which precision and training scheme will yield the lowest loss under specific compute constraints, helping to choose optimal hyperparameters.
24:45 FP4 Found Optimal for LLM Pre-training: Analysis, including real-world kernel speedups, shows that many practical LLM pre-training scenarios (e.g., for models like LLaMA 3, Qwen 2.5, Gemma 3) fall into regions where FP4, either for forward-only or full training, yields better accuracy under equivalent compute budgets.
28:17 Optimal Quantization Schemes Identified: The research identifies "Quest" rounding (standard deviation-based scaling to minimize quadratic error over Gaussian distribution) as the best for the forward pass and stochastic rounding for the backward pass, both utilizing Hadamard transforms for effective gradient normalization.
30:01 Crucial Role of Optimized CUDA Kernels: Achieving the predicted speedups necessitates highly optimized CUDA kernels to minimize the overhead of quantization-related operations and Hadamard transforms.
36:41 MXFP4 Data Format: The chosen 4-bit floating-point type is MXFP4 (E2M1 format), which uses FP8 (E8M0) for scaling factors, leveraging dedicated hardware support in Nvidia's Blackwell architecture.
40:50 Hadamard Transform Optimization: For small Hadamard sizes (e.g., 32), the operation is memory-bound, meaning its overhead can be effectively hidden by efficient data loading, allowing it to be integrated without significant performance penalties.
46:13 Fused Quantization Kernel: A custom Cutlass GEMM template is used for Hadamard transformation and quantization, featuring a custom epilogue that generates FP4 values, FP8 scales, and a clipping mask "almost for free" due to local computation and lack of synchronization.
50:02 Performance Breakdown and Speedups: For large matrices, the matrix multiplication dominates runtime, allowing quantization overhead to be hidden. Optimizations, such as larger thread block tiles (e.g., 128x128) and Triton-based rearrangement kernels, further reduce overhead. The current end-to-end speedups (MXFP4 vs. FP8) are approximately 2.2x for the forward pass and 1.5x for the backward pass, with ongoing optimization efforts.
56:30 Baseline Performance and Consistency: The proposed composite method (Quest + stochastic rounding) consistently outperforms existing 4-bit training baselines. Validation loss is demonstrated to be consistent with zero-shot performance on various downstream tasks, validating the robustness of their comparison methodology.
59:01 Conclusion and Future Outlook: The research provides a robust, compute-normalized method for comparing training methods and demonstrates that native FP4 training can be optimal for LLM pre-training, enabling better resource utilization.
1:07:33 Rationale for MXFP4: The choice of MXFP4 over NVFP4 was primarily due to its earlier standardization, longer-existing specification, and broader support across various accelerators (e.g., Chinese, Amazon), offering more "future-proof" compatibility. Performance differences between the two FP4 formats are expected to be minimal.