大量小算子

场景解析

Profiling数据发现Loss Scale在梯度更新引入大量Add、Mul小算子。

调优思路

融合小算子、并替换亲和API减少梯度更新阶段的add、mul算子调用数量。

优化方法

  1. 开启混合精度,并开启combine_grad=True开关。

    该环境变量在使用混合精度时,已默认提醒用户开启。

  2. 替换亲和API,函数说明请参考亲和库

样例参考

替换亲和函数。原始代码如下:

optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)  
model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0)

优化后代码如下:

optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=args.lr, momentum=args.momentum)  #NpuFusedSGD为亲和API函数
model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0, combine_grad=True)