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 软件生态,提升国产 GPU 在 AI 训练与推理领域的应用能力。
