推荐系统大多是从大规模CPU分布式计算开始的，这是由于早期的深度模型来源于简单的LR和FM的演进，如WDL，它的深度部分较为简单，整体还是IO较重，在GPU上的加速提升有限，而大规模GPU的成本又很高。随着序列拉长，Transformer结构的应用，模型变深变宽，复杂地主网络结构的应用，切换GPU计算能大幅地提升训练效率和推理性能，提升ROI。但是，简单地把CPU或消费级GPU上实现的代码搬到专业级GPU上，可能没有充分压榨到GPU的每一寸电路的性能。
混合精度，开启tensor core的利用
高性能专业级GPU的关键是利用好Tensor Core。Tensor Core 是 NVIDIA 在 Volta（V100）、Turing、Ampere（A100）、Hopper（H100） 及后续架构中引入的专用计算单元，用于加速深度学习和矩阵计算。其中，半精度（FP16）是 Tensor Core 计算的关键优化点。 Ampere（A100）之后，加入 TF32（TensorFloat-32），提供 FP32 级别的精度，但计算速度接近 FP16。
先简单介绍一下每种数据格式：
数据格式 | 简介 | 结构 (符号位+指数位+尾数位) | TensorFlow 版本支持 | GPU/硬件支持 | 加速情况（理论/实际） | 优点 | 缺点 |
FP32 | 标准单精度浮点数，适用于大多数高精度计算 | 1 + 8 + 23 | 所有版本均支持 | 所有 GPU 均支持（通用标准） | 基准运算，无专门低精度加速 | 精度高，数值稳定 | 计算量大、内存占用高，吞吐率低 |
TF32 | 改进型 FP32，在保持动态范围的同时利用 Tensor Core 加速 | 1 + 8 + 10 | TF 2.4 及以上 | NVIDIA Ampere（A100、H100）及以上 | 理论加速 2～4 倍（矩阵乘法、卷积等关键运算加速） | 保留 FP32 的动态范围，同时利用 Tensor Core 加速 | 精度低于 FP32，对细节要求较高的计算可能不够精确 |
FP16 | 半精度浮点数，数据量小、传输快，适合推理及部分训练 | 1 + 5 + 10 | TF 1.6 及以上 | NVIDIA Volta（V100）及以上 | 理论可达 4～8 倍；实际 2～4 倍左右（通常需 Loss Scaling 处理） | 内存占用小、数据传输快，适合推理 | 动态范围小，训练时易出现下溢问题，需要 Loss Scaling |
BF16 (TF16) | 16 位数据格式，保留 FP32 指数范围，训练更稳定，无需 Loss Scaling | 1 + 8 + 7 | TF 2.1 及以上 | TPU、部分支持 BF16 的 NVIDIA Ampere GPU（A100）及部分 CPU | 加速优势与 FP16 类似，实际中通常也能实现约 2～4 倍加速 | 数值范围大（与 FP32 相同），训练稳定性好，无需 Loss Scaling | 尾数位较少，精度低于 FP16，在部分精度敏感场景下可能不足 |
再直观感受下不同的数据格式的GPU加速情况：
在GPU预热后，执行4096x4096的矩阵乘法1k轮，不同精度下的耗时。本文提到的专业级GPU特指A100及以上的型号，V100仅拿来对比。
ㅤ | FP32 | TF32 | FP16 | BF16 |
V100 | 104.000ms | 不支持（104.017ms） | 15.961ms | 不支持（133.941ms） |
A100 | 72.219ms | 11.374ms | 5.666ms | 5.523ms |
可以看到，TF32/FP16/BF6利用了tensor core加速，因为进度差异差不多是两倍的耗时，FP32无法利用tensor core，耗时是半精度的十倍起。A100 的 Tensor Core 对 FP16 进行了大幅优化，使其计算速度远远快于 FP32。
- FP16 Tensor Core 吞吐量更高（比 FP32 高 8-16 倍）。
- FP32 只能在 CUDA Core 或 TF32 Tensor Core 上运行，而 FP16 直接在 Tensor Core 上计算，效率更高。
数据类型 | 计算方式 | 吞吐量（TFlops） |
FP32 | CUDA Core | 19.5 TFLOPS |
TF32（Tensor Core） | Tensor Core | 156 TFLOPS（8× FP32） |
FP16（Tensor Core） | Tensor Core | 312 TFLOPS（16× FP32） |
推荐系统切换半精度
工业界大多数的推荐模型基于tensorflow框架，且固化到了1.x的版本。尤其是大厂，喜欢拿来改改适配自家的生态，有个专门的组做性能优化，这就又让框架和开源版本岔开，牢牢地焊死在1.x的版本上。所以在大规模分布式CPU或小GPU上，TF的代码是以FP32全精度书写和维护的，全精度没有充分利用tensor core。
切换到FP16半精度并不是一个简单的问题，需要对老汤的模型代码做全面的改造和校验，比如处理好混合精度，避免数值下溢等问题，以及老汤模型的重新训练。
有一个可替代的方案是使用TF32. TF32 是 NVIDIA 在 Ampere 架构（如 A100）上引入的一种计算模式，用于加速 FP32 矩阵运算，矩阵乘法运算加速 2～4 倍，并且它和FP32的指数设置相同，只是削减了尾数，没有数值下溢出的风险。
以及思路相同的方案，使用BF16. 它的结构是1 位符号位 + 8 位指数位 + 7 位尾数位，与FP32的指数位相同，不易出现溢出/下溢问题，适用于深度学习训练（可避免 Loss Scaling），特别是 梯度累积、权重存储等。
但是参考上面的表格，在tf 1.x的版本上，这两种方案都没有有效的支持。
老老实实地回到半精度改造的路子上，核心是配合混合精度，避免数值下溢。
容易出现数值溢出的问题的地方：
- 在梯度累加、求和、均值计算等需要精度保证的操作中，建议采用FP32进行累积计算，避免因FP16精度不足而出现舍入误差累积。
- 像softmax、log、exp、以及Batch Normalization等对数值敏感的函数，通常应在FP32精度下计算，确保在指数运算和归一化过程中不出现下溢或上溢问题。
- 损失函数（如交叉熵、对数似然损失等）最好在FP32下执行，或使用Loss Scaling技术，在反向传播时对梯度进行缩放，以防止梯度过小而导致数值下溢。
- 除法运算以及矩阵求逆这类操作对精度要求较高，直接在FP16下计算可能不稳定，可以先转换为FP32计算后再转换回FP16。
在TF1.x的设定下，由于不支持TF32，许多敏感操作（例如梯度累加）仍在 FP32 下执行，而大部分计算密集型任务（如前向和反向传播中的卷积和全连接层）可以利用低精度加速。TF32 自动对这些关键运算进行加速，不需要开发者手动切换数据类型或实现额外的优化代码。所以即便手动做了混合精度，也应该升级TF版本或者改为pytorch，享受代码的简化和GPU极致的高性能。
最大的计算量还是出现在涉及到矩阵乘法的部分，它只要改成fp16的输入即可，但是对于大厂自行封装的组件来讲，为了确认它的逻辑正确性，通常要花更长的时间。
- 通过tensorboard观察FP32和FP16下各层梯度和激活的最大值、最小值、均值和标准差，如果发现某层数值异常（例如出现 NaN、Inf 或极端数值），可能就存在溢出或下溢风险。
- 逐层检验：针对模型中的关键层（如卷积层、全连接层和归约操作），分别用 FP32 和 FP16 进行前向计算，并对比输出。可以通过“局部测试”的方式逐层验证，确定哪些层在低精度下存在显著误差。
解决了半精度的问题，开启了对tensor core的利用，事情还没完，你很容易发现tensor core的利用率上不去，甚至只有个位数的利用率。
这个是由于推荐系统有非常多的特征，他们的处理涉及到大量非矩阵运算的算子，导致tensor core没法被数据填满，利用率低下。这时应该尽可能保证运算中大量地使用矩阵运算占大头的算子，不同的模型有不同的改造思路，但是减少繁杂地特征操作，利用多层的模型结构，以巨大的参数量完成拟合，是通用的方向。
把参数维度设置成64的倍数吧
混合精度正确地设置了只是第一步，这只是开启tensor core使用的条件。要关注tensor core是否被高效地利用，一些参数设置也是关键，在CPU下的线性思维和GPU上的并行计算思维是差异很大。
先做一个简单的实验，还是上面的测试程序，只继续在A100上将矩阵的维度改为4095.
ㅤ | FP32 | TF32 | FP16 | BF16 |
A100 （4096） | 72.219ms | 11.374ms | 5.666ms | 5.523ms |
A100 （4095） | 72.408ms（+0.2%） | 18.669ms（+64%） | 16.462ms（+190%） | 16.457ms（+197%） |
按非GPU的思维，当维度下降，计算量下降了，虽然下降的不多，但总不至于耗时大幅度上涨。这在FP32上基本上耗时一致，而对于tensor core可以优化的精度，耗时大幅度地上涨了。
这个现象的本质原因，是 Tensor Core（以及底层 GEMM kernel）的“切块”与“对齐”机制导致的：
- 切块（tiling）大小不匹配
- 不同数值格式的 fallback 差别
- 内存对齐与访存效率
知道了这个原理，首先维度最好设置成16的倍数，因为底层访存和矩阵库的设置，要设置成32（A100）/64（H100）的倍数。
推荐系统里，不按照这个规范来的比比皆是，举个最普遍的，所有特征embedding的concat之后的结果，就未必是64的倍数。那么你第一层的矩阵运算，就无法用满tensorcore的性能。
我最近有一个感悟，大厂应该想办法去拥抱开源，把生态打通，而不是基于一个开源的版本做闭源的改造，短期内可能有些收益，但是拉长了时间，闭源维护带来的人力成本，以及脱离了生态带来的长期负向损失是不可估量的。我记得docker调度管理刚开始的时候，k8s还不是最流行的，阿里的一个P8选了一个最简单的开源项目，他的理由是像我们阿里这种大厂，必定是要自己改造的，所以基座越简单越好。我现在常常把这个当笑话来讲，而算法领域的笑话就是tf的1.x版本。