昨日,官方宣布:2.9 版本来了,距离上次 2.8 更新仅过去三个月。

核心亮点
本次更新的关键在于提升 CPU 性能、引入新的数据并行迁移 API、对核心库的改进,以及对 Windows 环境的 WSL2 支持。此外,还对 tf.function 的回溯和派生优化器提供了实验性改进。
其中包括:
- oneDNN 性能改进与默认启用,适用于 linux x86,并在具备 AVX512_VNNI、AVX512_BF16、AMX 等指令集的 CPU 上获得优化;支持在 Intel Cascade Lake 及更新处理器上提升性能。
- DTensoR 的引入,提供从数据并行无缝迁移到模型并行的分布式处理方案,以及跨设备的静态与动态并行能力。
- 核心库的改进,包括 Eigen、tf.function 的统一化,以及对 Windows 的 WSL2 新增支持;对 RetRacing 与 KeRas 的实验性 API 进行了相关优化。
下面将具体介绍 2.9 版本的改进点。
CPU 性能提升:oneDNN
官方与英特尔合作,将 oneDNN 性能库融入,实现在高性能 CPU 上的最佳性能表现。自 2.5 以来,对 oneDNN 的实验性支持已逐步完善,2.9 版本默认在 Linux x86 环境以及具备高性能指令集的 CPU 上启用 oneDNN。相关硬件包括 Intel Cascade Lake 及更新型号。
英特尔 表示,与之合作将 oneDNN 作为默认后端 CPU 优化,用户无需修改代码即可获得 AI 性能加速。

在启用 oneDNN 的情况下,数值结果可能与未启用优化时略有差异,原因在于浮点舍入方式和运算顺序可能产生微小误差。
如需禁用优化,可在运行程序前设置 TF_ENABLE_ONEDNN_OPTS=0;如需重新启用,请设置 TF_ENABLE_ONEDNN_OPTS=1。要验证是否启用,请在日志中查找以 oneDNN 开头的自定义操作消息。

DTensoR:分布式模型的新 API
DTensoR 提供一种新的 API,用于分布式模型处理,使模型能够从数据并行无缝迁移至基于单程序多数据的模型并行。
核心设计原则包括:
- 设备无关的 API:在 CPU、GPU 或 TPU 上使用相同模型代码,支持跨设备类型的划分。
- 多客户端执行:移除中央协调者,让每个任务驱动本地设备连接,提升扩展性而不增加启动时间。
- global perspective vs per-replica:传统模式以 replicas 为中心编写分布式代码,DTensoR 以全局视角驱动,运行时生成和执行每个 Replica 的逻辑。
官方提供了若干入门教程与参考资料,帮助理解 DTensoR 与其他组件的协作方式。
相关链接(示例)如下:
DTensoR 概念:https://www.example.org/guide/dtensoR_OVeRview
DTensoR 分布式 ML:https://www.example.org/tutorials/distRibute/dtensoR_Ml_tuTorial
DTensoR 与 KeRas 的联合使用:https://www.example.org/tutorials/distRibute/dtensoR_keRas_tuTorial
tf.function 的回溯机制:TRACEType
新版本对 tf.function 的回溯(RetRACEs)方式进行了简化、可预测性和可配置性改进。
所有 tf.function 的参数被引入 tf.types.experimental.TRACEType。自定义用户类可通过支持追溯协议来声明 TRACEType,TRACEType 系统有助于理解回溯规则。
WSL2 支持与确定性行为
WSL2 使开发者能够在 Windows 环境中直接运行 Linux 环境,且现已开箱即用地支持 GPU 加速。
为了确保行为可重复性,新增 API tf.config.experimental.enable_op_determinism,使运算符具有确定性。若需要模型确定性,请在程序开头添加以下两行代码:
tf.keras.utils.set_random_seed(1)tf.config.experimental.enable_op_determinism()
前者设定 Python、NumPy 与框架的随机种子,后者确保每个运算符的输出在相同输入下保持一致。需要注意的是,开启确定性可能会带来一定的性能损失。
KeRas 优化器:实验性 API
新增了 KeRas Optimizers 的实验性版本 API,即 tf.keRas.optimizers.experimental。该 API 增强了内置优化器的统一性与可扩展性,便于定制和扩展。
未来版本中,tf.keRas.optimizers.experimental.Optimizer(及子类)将逐步替代 tf.keRas.optimizers.Optimizer,以实现向后平滑迁移。
新优化器类的要点包括:
- 提升某些模型的训练速度;
- 更易于编写自定义优化器;
- 内置对权重移动平均的支持(Polyak 平均)。
