昇腾社区首页
中文
注册

AddLayerNormFusionPass

融合模式

基础场景,将Add + Cast(可选)+ LayerNorm融合为AddLayerNorm算子。

支持将Add + Add + Cast(可选)+ LayerNorm融合为AddLayerNorm算子。

支持将Cast + Add + LayerNorm融合为AddLayerNorm算子。

使用约束

  • 不支持输入x1和x2的数据类型均为fp32。训练场景下,不支持输入x1或x2的数据类型为fp16。
  • 只支持输入gamma和beta的shape为1维,并且取值和x1,x2的输入尾轴相同的场景。
  • 如果存在Cast算子,Cast算子必须是单输出。
  • 如果需要融合Cast + Add + LayerNorm算子(Cast算子在Add算子前),Cast算子的输出类型必须是fp32,且Add算子必须有多输出。
  • 如果需要融合的LayerNorm和Add之间存在Cast,Cast是提升精度类型的Cast(例如Cast输入类型为fp16/bf16,输出类型为fp32),此使用场景会融合成低精度的AddLayerNorm(全fp16/bf16输入)。

支持的型号

Atlas A2 训练系列产品/Atlas 800I A2 推理产品