NPU亲和优化器替换
基本原理
当前昇腾AI处理器支持NPU亲和的融合优化器功能，针对表1中的优化器做了昇腾AI处理器亲和性优化，使性能提升。
- 前向传播。
- Loss乘缩放系数s。
- 反向传播。
- 梯度和参数更新，梯度和参数除以缩放系数s。
以上流程每个迭代执行一次，对梯度和参数的更新计算会包含连续多个小算子（如add、mul、sqrt等）下发，由于小算子在NPU上计算快，导致算子在CPU上的下发成为主要瓶颈。
操作步骤
优化器名称 |
样例原代码 |
修改后代码 |
|---|---|---|
sgd |
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr) |
optimizer = torch_npu.optim.NpuFusedSGD(model.parameters(), lr=args.lr) |
adadelta |
optimizer = torch.optim.Adadelta(model.parameters(), lr=args.lr) |
optimizer = torch_npu.optim.NpuFusedAdadelta(model.parameters(), lr=args.lr) |
lamb |
optimizer = torch.optim.Lamb(model.parameters(), lr=args.lr) |
optimizer = torch_npu.optim.NpuFusedLamb(model.parameters(), lr=args.lr) |
adam |
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) |
optimizer = torch_npu.optim.NpuFusedAdam(model.parameters(), lr=args.lr) |
adamw |
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) |
optimizer = torch_npu.optim.NpuFusedAdamW(model.parameters(), lr=args.lr) |
rmsprop |
optimizer = torch.optim.RMSprop(model.parameters(), lr=args.lr) |
optimizer = torch_npu.optim.NpuFusedRMSprop(model.parameters(), lr=args.lr) |
优化策略
使用昇腾AI处理器亲和的融合优化器替换原始优化器，如使用torch_npu.optim.NpuFusedRMSprop替换torch.optim.RMSprop。替换后重建训练的Profiling数据如下。使用案例请参见NPU亲和优化器替换调优案例。
结果验证
若用户需要验证具体的优化效果，可参考E2E Profiling数据采集与分析，在Chrome浏览器中打开msprof.json文件，查看梯度更新部分，如下。
通过上图可发现在梯度更新过程中，存在大量小算子下发。
对比发现算子数量明显减少。
