算子二进制调优

基本原理

PyTorch框架提供与算子编译相关的二进制配置参数,可设置模型编译时是否优先在线编译,以此优化模型训练性能。参数配置代码如下,可设置True或False。当前针对Atlas 训练系列产品,若不添加以下代码则默认为True。

torch_npu.npu.set_compile_mode(jit_compile=False)

用户可根据模型是固定shape或动态shape配置该参数以提高性能。

使用场景

用户在模型训练后,可参见PyTorch Analyse迁移分析工具,使用“dynamic_shape”模式,添加动态分析代码并运行脚本,根据分析结果判断为固定shape或动态shape,然后选择如下配置:

若将开关设置为False,则需安装二进制算子包,安装请参考CANN 软件安装指南常用操作 > 安装、升级和卸载二进制算子包章节。

操作步骤

  1. 动态shape场景下,在模型脚本的main_worker函数中配置进程级别的开关,配置为False。

    torch_npu.npu.set_compile_mode(jit_compile=False)

  2. 配置位置根据不同的训练拉起方式存在差异。

    • 单卡训练。正常拉起方式需将代码使能在main函数开始位置,mp.spawn方式拉起需配置在main_worker函数中,保证全部进程拉起时配置生效。
      • 正常拉起:
        if __name__ == '__main__':
            torch_npu.npu.set_compile_mode(jit_compile=False)
            main()
      • mp.spawn方式:
            mp.spawn(main_worker,...)
        ...
        def main_worker():
            torch_npu.npu.set_compile_mode(jit_compile=False)
    • 多卡训练。shell脚本、Python方式与单卡正常拉起方式配置相同,mp.spawn方式拉起需配置在main_worker函数中,保证全部进程拉起时配置生效。
      • shell脚本、Python方式:
        if __name__ == '__main__':
            torch_npu.npu.set_compile_mode(jit_compile=False)
            main()
      • mp.spawn方式:
            mp.spawn(main_worker,...)
        ...
        def main_worker():
            torch_npu.npu.set_compile_mode(jit_compile=False)