APEX

功能说明

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")

model, optimizer = amp.initialize(model, optimizer, opt_level="O2")

model, optimizer = amp.initialize(model, optimizer, opt_level="O3")

model, optimizer = amp.initialize(model, optimizer, opt_level="O2",loss_scale=128.0)

model, optimizer = amp.initialize(model, optimizer, opt_level="O2",loss_scale="dynamic")

推荐优先使用opt_level='O2', loss_scale=128.0的配置进行amp.initialize。若无法收敛，推荐使用opt_level='O1', loss_scale=128.0的配置进行amp.initialize。若依然无法收敛，推荐使用opt_level='O1', loss_scale=None的配置进行amp.initialize。

约束与限制

1. APEX当前版本的实现方式为python实现，不支持APEX中的自定义优化CUDA Kernel。
2. APEX当前版本只支持适配昇腾AI处理器的混合精度计算和多种融合优化器功能，其他功能暂未支持。融合优化器与原生优化器算法保持一致，但运算速度更快。

安装方法

请用户参考《CANN 软件安装指南》中“安装PyTorch”章节内的“安装APEX混合精度模块”，安装PyTorch框架对应的APEX即可。

使用说明

• 单卡训练场景。
  1. 导入混合精度模块。

     from apex import amp

  2. 在模型和优化器定义之后初始化APEX模块。

     model = ...
     optimizer = ...
     model, optimizer = amp.initialize(model, optimizer, combine_grad=True)

  3. 改写梯度反向传播loss.backward()。

     loss = criterion(…) 
     #将loss.backward()替换为如下形式
     with amp.scale_loss(loss, optimizer) as scaled_loss:     
         scaled_loss.backward() 
     optimizer.step()

• 多卡训练场景。
  • （推荐）使用APEX混合精度并关闭combine_ddp开关，并开启PyTorch框架自带的DistributedDataParallel（DDP）模式，DDP模式可参考PyTorch官方文档使用。

    from torch.nn.parallel import DistributedDataParallel as DDP
    from apex import amp
    ...
    model = ...
    optimizer = ...
    #在模型、优化器定义之后，初始化APEX模块。
    model, optimizer = amp.initialize(model, optimizer, combine_ddp=False)

  • 使用APEX混合精度并开启combine_ddp开关，并关闭PyTorch框架自带的DDP模式。

    from apex import amp
    ...
    model = ...
    optimizer = ...
    #在模型、优化器定义之后，初始化APEX模块。
    model, optimizer = amp.initialize(model, optimizer, combine_ddp=True)

更多混合精度模块的使用可参见官方文档。

APEX特性支持

• 替换亲和优化器函数。

  适配后的APEX针对adadelta/adam/sgd/lamb做了昇腾AI处理器亲和性优化，得到的NPU融合优化器与原生算法保持一致，但运算速度更快。使用时只需将原有优化器替换为apex.optimizers.***，***为优化器名称，例如NpuFusedSGD。

  样例原代码：

  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)    
  model, optimizer = amp.initialize(model, optimizer, opt_level='O2', loss_scale=32.0, combine_grad=True)

• APEX使用小技巧：若希望通过自定义设置部分算子为float16提升性能或设置部分算子为float32提升精度，可以通过如下方法实现。

  amp.register_half_function(torch, 'bmm') #填入算子API名称。以bmm为例，bmm会强制使用half进行计算 
  amp.register_float_function(torch, 'bmm') #填入算子API名称。以bmm为例，bmm会强制使用float进行计算

  

  当optimizer的可更新参数分为不同的多组，每组使用不同的策略，这种场景下，融合优化器只能优化掉第一个model分组的参数。因此，当优化器的参数相同时，尽量合并分组。

  优化前：

  optimizer = apex.optimizers.NpuFusedSGD([
               {'params': model.sharedNet.parameters()},
               {'params': model.bottleneck.parameters()},
               {'params': model.domain_classifier.parameters()},
               {'params': model.dcis.parameters()},
               {'params': model.source_fc.parameters(), 'lr': LEARNING_RATE},
           ], lr=LEARNING_RATE / 10, momentum=args.momentum, weight_decay=args.l2_decay)

  优化后：

  optimizer = apex.optimizers.NpuFusedSGD([
               {'params': list(model.sharedNet.parameters()) + list(model.bottleneck.parameters()) +  list(model.domain_classifier.parameters())+list(model.dcis.parameters())},
               {'params': model.source_fc.parameters(), 'lr': LEARNING_RATE},
           ], lr=LEARNING_RATE / 10, momentum=args.momentum, weight_decay=args.l2_decay)

• 开关分布式训练性能。

  适配后的APEX针对数据并行场景做了昇腾AI处理器亲和性优化，支持利用融合grad进行加速，同时保持计算逻辑一致性。通过开启combine_ddp开关，也就是将amp.initialize()接口参数combine_ddp设置为True并关闭框架的DDP模式，即可开启该功能。

  model, optimizer = amp.initialize(model, optimizer, combine_ddp=True) #配置运算加速参数combine_ddp

  

  • 当前版本的实现方式主要为python实现，不支持AscendCL或者CUDA优化。
  • 当前昇腾AI设备暂不支持原始APEX的FusedLayerNorm接口模块，如果模型原始脚本文件使用了FusedLayerNorm接口模块，需要在模型迁移过程中将脚本头文件from apex.normalization import FusedLayerNorm替换为from torch.nn import LayerNorm。
  • 如果在DDP初始化后更改优化器，会导致主卡生效，副卡不生效，从而使模型出错。为了规避这个问题，需要将优化器更改放到DDP初始化前。