【CNMO 科技消息】8 月 21 日消息,摩尔线程正式发布 TensorFlow-MUSA v1.5.0。作为面向 MUSA 统一系统架构 GPU 推出的 TensorFlow 插件,该版本进一步完善了对搜索、广告、推荐(以下简称 " 搜广推 ")业务的支持,开发者可以在摩尔线程全功能 GPU 上直接开展 AI 模型训练与推理,无需对原有 TensorFlow 模型代码进行大幅修改。此次更新还针对搜广推业务的模型特点进行了专项优化,并已在 GitHub 开源,同时提供预构建镜像,方便开发者快速部署。

搜广推业务通常需要处理海量用户行为数据和大规模稀疏特征,其计算特点与视觉、语音等常见 AI 任务存在较大区别。一方面,Embedding 参数规模较大,模型结构也需要频繁迭代;另一方面,训练阶段需要较高的吞吐能力和分布式计算能力,而在线推理又对延迟和吞吐提出较高要求。由于训练和推理往往需要共用模型代码及工程链路,因此 AI 框架不仅需要具备较为完整的算子支持,还需要兼顾混合精度和分布式计算能力。
针对这些需求,TensorFlow-MUSA v1.5.0 在框架层面实现了对 TensorFlow 数据流图、自动微分以及分布式训练等核心机制的支持,并将 MUSA 设备作为 TensorFlow 原生设备进行注册。对于原本基于 TensorFlow 开发的搜广推模型而言,可以继续沿用原有开发方式,将模型迁移至摩尔线程全功能 GPU 上运行,从而降低适配过程中需要进行的代码和工程调整。
在模型兼容性方面,TensorFlow-MUSA v1.5 覆盖范围进一步扩大。该版本不仅继续支持 ResNet 等经典视觉、语音模型,还加入了大量搜广推模型,包括 Wide&Deep、FNN、PNN、DeepFM、NFM、MLR、SharedBottom 等经典模型,以及 DCN、xDeepFM、AutoInt、FGCNN、FiBiNet、FWFM 等特征交叉模型。同时,针对用户行为序列建模和多任务学习场景,该版本支持 DIN、DIEN、DSIN、BST、MMoE、ESMM、PLE 等模型。
此外,TensorFlow-MUSA v1.5 还针对近年来出现的新型搜广推架构进行了适配,支持 EDCN、DCNmix、Wukong、RankMixer、OneTrans 和 TokenMixerLarge 等模型。摩尔线程表示,这些模型覆盖了从传统推荐算法到新型模型架构的多个应用场景,可以满足开发者在不同搜广推业务中的训练和推理需求。
软件版本方面,TensorFlow-MUSA v1.5 支持 TensorFlow 2.6.1 和 2.15.1,并可配合 MUSA SDK 4.3.5 或 5.1.0 使用。摩尔线程表示,考虑到 TensorFlow 2.15.1 之后的功能变化相对有限,同时部分早期特性被移除,后续将主要维护 TensorFlow 2.15.1 版本,并保持前向兼容,同时根据开发者需求评估其他版本的支持。
性能方面,摩尔线程公布的实验室测试数据显示,基于 AI 训推一体智算卡 MTT S5000 运行搜广推模型时,在 BF16、TF32 等数据精度以及多流执行、算子融合、图优化和 XLA 编译优化等能力支持下,部分模型表现出较明显的性能提升。以 TokenMixerLarge 为例,单步训练耗时为 138.048ms,相较国际主流 GPU 的加速比达到 1.81 倍。摩尔线程同时给出的整体测试结果显示,部分前沿搜广推模型训练吞吐提升幅度达到 1.5 倍至 1.85 倍。不过,官方也注明,相关数据来自摩尔线程实验室,实际性能会受到模型实现、硬件配置、软件版本以及训练参数等因素影响。

目前,TensorFlow-MUSA v1.5.0 已经在 GitHub 开源,开发者可以按照官方说明自行编译安装,也可以直接使用预构建镜像。此次版本更新将重点从通用模型支持进一步延伸至搜广推等高计算需求场景,摩尔线程也在通过框架适配、模型覆盖和性能优化等方式,完善 MUSA GPU 的软件生态。