使用自定义融合算子
原理
由于有些算子在device侧NPU上计算过快,而导致device侧需要等待HOST侧CPU下发任务进行计算。此等待过程就产生了性能瓶颈。自定义融合算子的更多内容请参见《图融合和UB融合规则参考》。
问题定位
- 参见Profiling数据采集及分析在训练脚本中使能profiling。
- 对照表1查看profiling数据中是否存在连续的小算子且耗时占用较多,对应到相应的代码行并进行优化。
小算子/算子 |
自定义融合算子/自定义算子 |
对应亲和库 |
说明 |
---|---|---|---|
transpose |
matmul_transpose() |
class MatmulApply() |
小算子融合。 |
matmul |
|||
add |
fuse_add_softmax_dropout() |
def fuse_add_softmax_dropout() |
小算子融合。 |
softmax |
|||
dropout |
|||
transpose |
npu_confusion_transpose() |
- |
小算子融合。 |
reshape |
|||
sigmoid |
npu_silu() |
class SiLU() |
Silu是Swish的一种特例。 |
mul |
亲和库详细内容请参见API列表中“{对应的PyTorch版本号}>PyTorch Adapter自定义API>torch_npu.contrib”章节。
优化策略
将存在下发瓶颈的算子替换成自定义融合算子。如将transpose+matmul融合为matmul_transpose,add+softmax+dropout融合为npu_dropout_with_add_softmax,transpose+reshape融合为torch.npu_confusion_transpose。调优案例请参见自定义融合算子替换优化调优案例。
父主题: 算子下发性能优化