探索4比特浮点训练的最优设计组合,提升效率与稳定性。
Elucidating the Design Space of FP4 training
- 基于量化梯度框架,系统分析不同稳定化方法的计算开销。
- 发现哈达玛变换+张量缩放+随机舍入组合效果最佳。
- 建议使用UE5M3作为缩放因子,平衡范围与精度代价。
基础模型日益增长的计算需求推动了低精度训练研究,4比特浮点(FP4)格式成为提升硬件吞吐量的前沿方向。尽管已有多种技术用于稳定FP4训练,但它们往往孤立存在,且计算开销差异大、不明确。本文旨在统一理解FP4训练的设计空间,提出一个基于量化梯度的微尺度量化框架,支持对前向与反向传播中不同稳定化方法的理论分析。基于该框架构建模拟器,在回归、图像分类、扩散模型和语言模型等广泛任务上开展大规模实证研究,系统评估数千种技术组合(包括新型梯度近似、舍入策略和缩放方法),识别出性能与开销最优权衡的配置。结果表明,结合哈达玛变换、张量缩放与随机舍入的方案表现最佳;进一步发现,采用UE5M3作为缩放因子可在范围与精度间取得良好平衡,且计算开销可控。
原文摘要 · Abstract (English)
The increasing computational demands of foundation models have spurred research into low-precision training, with 4-bit floating-point (\texttt{FP4}) formats emerging as a frontier for maximizing hardware throughput. While numerous techniques have been proposed to stabilize \texttt{FP4} training, they often present isolated solutions with varying, and not always clear, computational overheads. This paper aims to provide a unified view of the design space of \texttt{FP4} training. We introduce a comprehensive, quantisation gradient-based framework for microscaling quantization that allows for a theoretical analysis of the computational costs associated with different stabilization methods on both the forward and backward passes. Using a simulator built on this framework, we conduct an extensive empirical study across a wide range of machine learning tasks, including regression, image classification, diffusion models, and language models. By systematically evaluating thousands of combinations of techniques, such as novel gradient approximations, rounding strategies, and scaling methods, we identify which configurations offer the most favourable performance-to-overhead trade-off. We find that the techniques enabling the best trade-off involve carefully combining Hadamard transformations, tensor scaling and stochastic rounding. We further find that using \texttt{UE5M3} as a scaling factor potentially offers a good compromise between range and precision with manageable computational overhead.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。