导语自2015年TensorFlow开源以来，伴随着深度学习的迅猛发展，通用深度学习框架经历了10年的高速发展，大浪淘沙，余者寥寥。曾几何时，也有过性能与易用性之争，也有过学术界和工业界之分，但随着本轮大模型应用的推波助澜，PyTorch无疑已经成为事实上的大模型“标准框架”。时至今日，PyTorchAOTCompile特性的发布更是直接撕下了TensorFlow最后一块“易于部署”的遮羞布。社区活跃度、性能和易用性，数张无形的大手，推搡着我们去拥抱更加现代化的新质生产力——PyTorch。一、Torch做推荐，到底行不行？
基于上述背景，Gauss总结了实时推荐系统和生成式推荐对建模工具链的4大要求，并依此重新审视PyTorch框架，决策其应用于推荐场景的可能性。
关键技术1——动态词向量
大规模稀疏的词向量，是时下个性化推荐系统的基础配置，TensorFlow生态中存在TFRA[3]扩展，以便支持百亿量级动态变化的词向量。在PyTorch生态中，原不具备对应的功能，感谢TorchRec团队的贡献，其在PyTorch官方扩展TorchRec中引入了动态词向量的支持[4]，极大的降低了这部分的入门门槛。
关键技术2——高性能、强一致的推理实现
关键技术3——FlashAttention
关键技术4——容灾与弹性训练
实时推荐系统通常要求流式的进行模型训练，即便是分钟级别的训练中断也会对推荐效果产生不良影响[9]。此外，流式的样本管线通常存在波峰与波谷效应，如果训练资源一直保持波峰期的级别，将会造成极大的浪费。支持在不影响模型效果的前提下，在波峰波谷期动态的扩容与缩容，是评判一个实时训练框架是否完备的重要标准之一。额外的，低优（低价）GPU资源在公司内的推广使用，也为弹性训练打下了政策基础。框架方面，TensorFlow本身并不具备动态扩缩容的能力，但配合第三方框架（例如Horovod）可以实现简单的弹性训练。PyTorch内置了的能力，但具体应用到推荐领域的大规模稀疏模型上，也面临恢复续训和扩缩容速度较慢，不能秒级容错的局限。综上，在容灾与弹性方面，TensorFlow与PyTorch的能力均是捉襟见肘，需要额外补齐。
除了上述关键技术之外，分布式训练、异构设备（尤其是910B）的支持，对推荐模型的训练和推理也十分重要，PyTorch在这些方面与TensorFlow相比，各有千秋，总体上来说，均不会形成掣肘，这里就不过多展开了。
综上，今日之PyTorch不仅完全胜任推荐系统和生成式推荐的核心要求，与TensorFlow相比，在其中若干方面甚至更胜一筹。Gauss团队基于PyTorch生态快速搭建系统验证了上述想法，此外，还针对生态中的关键技术进行了一轮优化，在生成式推荐场景下，相比TensorFlow实现，训练速度提升3x倍，结合multi-targets[2]特性，推理吞吐提升数十倍。
二、Gauss-Torch工具链
图1Gauss-Torch模型训练推理示意
上图展示了Gauss-Torch在流式模型实时训练中的核心流程。在Embedding存储方面，Gauss-Torch沿用了Gauss-TensorFlow中TierPS与WePS[8]的双引擎设计，在训练时采用成本友好的TierPS，在推理时采用高可用的WePS（不上线的模型、非实时模型可以不部署）。与Gauss-TensorFlow不同的是，TierPS中去除了GPU缓存的部分，这部分作为Gauss-Torch的核心组件单独进行设计。训练过程中，框架会定期（秒级）向WePS同步有更新的Embedding参数，同时借助HDFS/WFS同步剩余模型参数，保障推理参数的实时性。
在Embedding查询部分，Gauss-Torch并没有直接采用TorchRec中的分布式缓存策略，因为其在多机多卡场景下的扩展性并不好（在通讯占比较高的场景下，2机性能约为单机的1.6x左右），而是借鉴发表在SOSP'23的工作Bagpipe[7]实现了一套流水线感知的单机Embedding缓存，通过启发式的预取与淘汰策略，实现了高通讯压力下线性的扩展能力。
自定义Attention融合算子方面，Gauss-Torch初期基于Triton实现了一版HSTU[2]专用融合算子，性能相比于使用Torch算子拼接而成的实现已有不错的提升，近期团队内使用更加底层的cuda+cutlass接口重构了这部分实现，单算子性能相较于Triton版本又有了2x~4x的提升。当然，性能只是一个方面，更重要的是赋能算法同学尝试不同的自定义Attention实现，快速实现业务价值。Gauss团队目前正基于Torch提供的编译优化技术，着手抽象更加灵活的Attention接口，届时，一天内开发高性能融合Attention核将不再是奢望。
推理方面，Gauss-Torch采用了社区最近几个版本发布的AOTInductor特性，其通过TorchDynamo将用户模型转化为一组顺序执行的代码，再使用TorchInductor将其编译为可执行的二进制程序发布上线，几乎完全兼容所有Torch算子和用户自定义算子。与Torch其他在线推理方案相比（ONNX，TensorRT）保障性能的同时，具备优秀的兼容性与训练/推理一致性。
最后，在容灾和弹性训练方面，Gauss-Torch基于Ray，以Torch-Elastic为基础，摒弃了必须从checkpoint恢复的机制，结合多副本的设计实现了秒级故障恢复、分钟级弹性扩缩容的能力。
三、结语
参考文献
[2]ZhaiJ,LiaoL,LiuX,etal.Actionsspeaklouderthanwords:Trillion-parametersequentialtransducersforgenerativerecommendations[J].arXivpreprintarXiv:2402.17152,2024.
[7]AgarwalS,YanC,ZhangZ,etal.Bagpipe:Acceleratingdeeprecommendationmodeltraining[C]//Proceedingsofthe29thSymposiumonOperatingSystemsPrinciples.2023:348-363.
[8]SimaC,FuY,SitMK,etal.Ekko:A{Large-Scale}deeplearningrecommendersystemwith{Low-Latency}modelupdate[C]//16thUSENIXSymposiumonOperatingSystemsDesignandImplementation(OSDI22).2022:821-839.
[9]HeX,PanJ,JinO,etal.Practicallessonsfrompredictingclicksonadsatfacebook[C]//ProceedingsoftheEighthInternationalWorkshoponDataMiningforOnlineAdvertising.2014:1-9.