MPX-JAX 的混合精度训练-计算机科学-神经网络训练-深度学习.pdfVIP

MPX:JAX的混合精度训练

11

AlexanderGräfeandSebastianTrimpe

1InstituteforDataScienceinMechanicalEngineering(DSME),RWTHAachenUniversity

混合精度训练近年来已成为提高神经网络训练效率不可或缺的工具。同时，JAX作为一

款多功能机器学习工具箱，也越来越受欢迎。然而，它目前缺乏对混合精度训练的强大支

持。我们提出了MPX，这是一个用于JAX的混合精度训练工具箱，它可以简化和加速大

规模神经网络的训练，同时保持模型准确性。MPX无缝集成到诸如Equinox和Flax等

流行工具箱中，允许用户只需进行少量修改即可将全精度管道转换为混合精度版本。通

过将输入和输出都转换为半精度，并引入动态损失缩放机制，MPX缓解了在半精度计算

中常见的梯度欠流和溢出问题。其设计继承了JAX类型提升行为中的关键特性，确保操

本作以正确的精度进行，并允许根据需要选择性地强制执行全精度（例如，求和、均值或

译softmax）。此外，MPX还提供了自动创建和管理混合精度梯度及优化器的封装程序，从而

中可以轻松集成到现有的JAX训练管道中。MPX的源代码、文档和使用示例可在github.

com/Data-Science-in-Mechanical-Engineering/mixed_precision_for_JAX获取。

2

v

2Githubmixed_precision_for_JAX

1

3

3

0

7.1介绍

0

5

2近年来，JAX(DeepMindetal.,2020)作为训练神经网络的流行工具崭露头角。JAX的一个突出特

:点是其即时编译功能，这显著加速了在GPU和TPU等硬件加速器上的代码执行速度。此外，通过

v

i

x用户友好且高效的分片机制，JAX支持跨计算集群的大规模多设备和多节点训练。

r

a为了扩展JAX的功能，已经开发了几个库，包括Equinox、Flax和Optax(DeepMindetal.,2020;

Heeketal.,2024;KidgerandGarcia,2021).尽管有了这些进步，JAX在支持混合精度训练方面

仍然存在不足，这是一种处理大型模型时非常重要的方法(Micikeviciusetal.,2017).虽然之前有

JMP(jmp)可用于此目的，但它已不再维护，并且无法处理现代神经网络工具箱中常用的任意JAX

PyTrees。

混合精度训练通过在半精度浮点格式（16bit）中执行前向和后向传递，在优化大型模型训练性能方

面发挥着关键作用。该技术显著减少了中间激活的内存需求，并通过最小化内存传输开销来加速训

练，同时利用专门的半精度张量核心（当硬件支持时），并且不会影响准确度(Micikeviciusetal.,

2017)。

为了简化在JAX中实现混合精度训练，我们引入了MPX，一个专为JAX设计的混合精度训练工具

箱。MPX基于JMP的实现，并利用Equinox的灵活性来克服处理任意PyTrees相关的限制。它还

引入了梯度变换和优化器封装，使得从全精度管道转换到混合精度的过程变得无缝且修改量最小。

MPX的源代码、文档和使用示例可在/Data-Science-in-Mechanical-Engineering/

mixed_precision_for_JAX获取。

©2025InstituteforDataScienceinMechanicalEngineering(DSME),RWTHAache