用8位浮点训练百万级标签分类模型,内存占用降为六分之一。
ELMO: Efficiency via Low-precision and Peak Memory Optimization in Large Output Spaces
- 全量使用Float8低精度训练,结合求和优化避免误差累积。
- 300万标签模型仅需6.6GiB显存,比当前最优方法节省33.1GiB。
- 适合大规模推荐与标签系统,显著降低部署门槛。
大规模输出空间(即极端多标签分类,XMC)常见于大规模标签系统与商品推荐,标签数可达数十万至百万。此时线性分类头虽仅占模型极小部分,却成为计算与内存的主要瓶颈。现有主流方法依赖混合精度训练(FP16-FP32),存在不稳定性且内存效率低。现有低精度方法通常仍保留分类层高精度。本文提出ELMO,一种纯低精度训练框架,采用BFloat16与Float8数据类型。通过引入Kahan求和与随机舍入,实现分类头完全在Float8下训练,无需单精度主权重或张量缩放。结合梯度融合与分块优化,显存显著降低:300万标签模型仅需6.6 GiB GPU内存,相较当前最优方法Renee的39.7 GiB减少约83%,且不损失准确率。
原文摘要 · Abstract (English)
Large output spaces, also referred to as Extreme multilabel classification (XMC), is a setting that arises, e.g., in large-scale tagging and product-to-product recommendation, and is characterized by the number of labels ranging from hundreds of thousands to millions. This means that the linear classification head, usually only a tiny fraction of the overall model, turns into the main driver for compute and memory demand. Current state-of-the-art XMC methods predominantly rely on FP16-FP32 mixed-precision training, which we show can be unstable, and inefficient in terms of memory usage and computational overhead. Meanwhile, existing low-precision methods typically retain higher precision for the classification layer. In this work, we propose ELMO, a pure low-precision training framework for XMC models using BFloat16 and Float8 data types. By leveraging Kahan summation and stochastic rounding, we demonstrate that XMC models can be effectively trained entirely in Float8, without relying on single-precision master weights or tensor scaling. Low-precision training, combined with our proposed memory optimizations -- gradient fusion and chunking -- enables significant reductions in GPU memory usage. For example, we train a 3-million-label XMC model with only 6.6 GiB of GPU memory, compared to the 39.7 GiB required by the optimized SOTA method, Renee without compromising accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。