Apex Patch

简介

Apex Patch以代码patch的形式发布，用户通过对原始Apex进行patch，可以在华为昇腾AI处理器上，使用Apex的自动混合精度训练功能进行模型训练，提升AI模型的训练效率，同时保持模型的精度和稳定性。此外，Apex-patch额外提供了如梯度融合、融合优化器等，以提升部分场景下模型在昇腾NPU上的训练效率，供用户选择使用。

版本说明

Apex-patch配套软件

CANN版本	PyTorch版本	Ascend Extension for PyTorch版本	Python版本	Apex 版本或代码分支
8.5.0	2.6.0	v2.6.0-7.3.0	Python3.9x,Python3.10x,Python3.11x,Python3.12x	master
8.5.0	2.7.1	v2.7.1-7.3.0	Python3.9x,Python3.10x,Python3.11x,Python3.12x	master
8.5.0	2.8.0	v2.8.0-7.3.0	Python3.9x,Python3.10x,Python3.11x,Python3.12x	master
8.5.0	2.9.0	v2.9.0-7.3.0	Python3.9x,Python3.10x,Python3.11x,Python3.12x	master
8.3.RC1	2.1.0	v2.1.0-7.2.0	Python3.8x,Python3.9x,Python3.10x,Python3.11x	master
8.3.RC1	2.6.0	v2.6.0-7.2.0	Python3.9x,Python3.10x,Python3.11x	master
8.3.RC1	2.7.1	v2.7.1-7.2.0	Python3.9x,Python3.10x,Python3.11x	master
8.3.RC1	2.8.0	v2.8.0-7.2.0	Python3.9x,Python3.10x,Python3.11x	master
8.2.RC1	2.1.0	v2.1.0-7.1.0	Python3.8x,Python3.9x,Python3.10x,Python3.11x	master
8.2.RC1	2.5.1	v2.5.1-7.1.0	Python3.9x,Python3.10x,Python3.11x	master
8.2.RC1	2.6.0	v2.6.0-7.1.0	Python3.9x,Python3.10x,Python3.11x	master
8.1.RC1	2.1.0	v2.1.0-7.0.0	Python3.8x,Python3.9x,Python3.10x,Python3.11x	master
8.1.RC1	2.3.1	v2.3.1-7.0.0	Python3.8x,Python3.9x,Python3.10x	master
8.1.RC1	2.4.0	v2.4.0-7.0.0	Python3.8x,Python3.9x,Python3.10x	master
8.1.RC1	2.5.1	v2.5.1-7.0.0	Python3.9x,Python3.10x,Python3.11x	master
8.0.0	2.1.0	v2.1.0-6.0.0	Python3.8x,Python3.9x,Python3.10x,Python3.11x	master
8.0.0	2.3.1	v2.3.1-6.0.0	Python3.8x,Python3.9x,Python3.10x	master
8.0.0	2.4.0	v2.4.0-6.0.0	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC3	2.1.0	v2.1.0-6.0.rc3	Python3.8x,Python3.9x,Python3.10x,Python3.11x	master
8.0.RC3	2.3.1	v2.3.1-6.0.rc3	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC3	2.4.0	v2.4.0-6.0.rc3	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC2	1.11.0	v1.11.0-6.0.rc2	Python3.7x(Python3.7.5及以上),Python3.8x,Python3.9x,Python3.10x	master
8.0.RC2	2.1.0	v2.1.0-6.0.rc2	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC2	2.2.0	v2.2.0-6.0.rc2	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC2	2.3.1	v2.3.1-6.0.rc2	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC1	1.11.0	v1.11.0-6.0.rc1	Python3.7x(Python3.7.5及以上),Python3.8x,Python3.9x,Python3.10x	master
8.0.RC1	2.1.0	v2.1.0-6.0.rc1	Python3.8x,Python3.9x,Python3.10x	master
8.0.RC1	2.2.0	v2.2.0-6.0.rc1	Python3.8x,Python3.9x,Python3.10x	master
7.0.0	1.11.0	v1.11.0-5.0.0	Python3.7x(Python3.7.5及以上),Python3.8x,Python3.9x,Python3.10x	master
7.0.0	2.0.1	v2.0.1-5.0.0	Python3.8x,Python3.9x,Python3.10x	master
7.0.0	2.1.0	v2.1.0-5.0.0	Python3.8x,Python3.9x,Python3.10x	master

环境部署

在安装Apex Patch之前，请参考Apex-patch配套软件，安装最新昇腾软件栈。

推荐使用docker，在容器内编译：参考链接。

建议用户以非root用户做环境的安装。若使用容器环境编译，建议使用普通用户，本仓库提供的Dockerfile仅供参考。请用户关注容器挂载目录安全性，避免系统路径，推荐只挂载业务路径，避免不必要的安全问题。

获取昇腾适配的Apex-patch源码

git clone -b master https://gitcode.com/Ascend/apex.git
cd apex/

编译apex的二进制包

1、请确保torch已安装，setuptools版本推荐使用65.7.0。若出现报错error: invalid command 'bdist_wheel'，须更新setuptools版本，如75.6.0及以上。请用户自行确保对应版本无重大漏洞。

2、执行如下命令（支持Python3.7-3.12，确保Python3.x命令存在）。

bash scripts/build.sh --python=3.7

过程中会自动拉取apex官方源码，请保证网络畅通，生成的二进制包在apex/dist目录下。

安装

进入apex/dist目录，执行以下命令：

cd apex/dist/
pip3 uninstall apex
pip3 install --upgrade apex-0.1+ascend-{version}.whl  # version为Python版本和CPU架构

如需要保存安装日志，可在pip3 install命令后面加上参数 --log <PATH>，并对您指定的目录<PATH>做好权限管控。

卸载

Pytorch框架训练环境的卸载可以参考昇腾官方文档。

Apex及Apex-patch的卸载只需执行命令：

pip3 uninstall apex

如需要保存卸载日志，可在pip3 install命令后面加上参数 --log <PATH>，并对您指定的目录<PATH>做好权限管控。

硬件配套

昇腾训练设备包含以下型号，都可作为PyTorch模型的训练环境，并使能Apex-patch相关功能。 | 产品系列 | 产品型号 |

|-----------------------|----------------------------------| | Atlas 训练系列产品 | Atlas 800 训练服务器（型号：9000） | | | Atlas 800 训练服务器（型号：9010） | | | Atlas 900 PoD（型号：9000） | | | Atlas 900T PoD Lite | | | Atlas 300T 训练卡（型号：9000） | | | Atlas 300T Pro 训练卡（型号：9000）| | Atlas A2 训练系列产品 | Atlas 800T A2 训练服务器 | | | Atlas 900 A2 PoD 集群基础单元 | | | Atlas 200T A2 Box16 异构子框 | | Atlas A3 训练系列产品 | Atlas 800T A3 训练服务器 | | | Atlas 900 A3 SuperPoD 超节点 | | | A200T A3 Box8 超节点服务器 |

快速上手

自动混合精度（AMP）

导入apex包，使用apex.amp.initialize方法包装模型和优化器并指定opt_level为O1，并使用apex.amp.scale_loss方法包装损失值和优化器对损失进行动态缩放，实现混合精度训练：

    import torch
+   import apex  # 导入apex包
    
    
    epoch = 10  # 训练轮次
    loss_function = torch.nn.MSELoss()  # 损失函数
    input_data = [torch.rand(16, device='npu') for i in range(10)]  # 构造简易模型输入
    dataset = [(data, 0.5 * torch.sum(data).unsqueeze(0)) for data in input_data]  # 构造简易数据集
    model = torch.nn.Linear(in_features=16, out_features=1, device='npu')  # 构造简易模型
    optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3)  # 创建优化器
+   # 使用apex.amp.initialize方法包装模型和优化器进行混合精度训练，opt_level取'O0'时，表示使用FP32训练，opt_level取'O1'或'O2'时，表示使用混合精度训练
+   model, optimizer = apex.amp.initialize(models=model, optimizers=optimizer, opt_level='O1')
    for i in range(epoch):
        for x, target in dataset:
            optimizer.zero_grad()
            output = model(x)
            loss = loss_function(output, target)
-           loss.backward()
+           # 使用apex.amp.scale_loss方法包装损失值和优化器进行动态损失缩放
+           with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
+               scaled_loss.backward()
            optimizer.step()

详细使用方式请请参考Apex官方文档。

使用融合梯度进行scale/unscale优化加速

导入apex包，使用apex.amp.initialize方法包装模型和优化器并将参数combine_grad设置为True，示例如下：

    import torch
+   import apex  # 导入apex包
    
    
    epoch = 10  # 训练轮次
    loss_function = torch.nn.MSELoss()  # 损失函数
    input_data = [torch.rand(16, device='npu') for i in range(10)]  # 构造简易模型输入
    dataset = [(data, 0.5 * torch.sum(data).unsqueeze(0)) for data in input_data]  # 构造简易数据集
    model = torch.nn.Linear(in_features=16, out_features=1, device='npu')  # 构造简易模型
    optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3)  # 创建优化器
+   # 使用apex.amp.initialize方法包装模型和优化器进行混合精度训练，opt_level='O1'表示使用混合精度训练，combine_grad=True表示开启融合梯度
+   model, optimizer = apex.amp.initialize(models=model, optimizers=optimizer, opt_level='O1', combine_grad=True)
    for i in range(epoch):
        for x, target in dataset:
            optimizer.zero_grad()
            output = model(x)
            loss = loss_function(output, target)
-           loss.backward()
+           # 使用apex.amp.scale_loss方法包装损失值和优化器进行动态损失缩放
+           with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
+               scaled_loss.backward()
            optimizer.step()

使用融合优化器优化加速

导入apex包，使用apex.amp.initialize方法包装模型和优化器并将参数combine_grad设置为True，并将torch原生优化器torch.optim.xxx替换为apex.optimizers.xxx, 其中xxx为融合优化器名称，示例如下所示，apex-patch支持的优化器见特性介绍。

    import torch
+   import apex  # 导入apex包
    
    
    epoch = 10  # 训练轮次
    loss_function = torch.nn.MSELoss()  # 损失函数
    input_data = [torch.rand(16, device='npu') for i in range(10)]  # 构造简易模型输入
    dataset = [(data, 0.5 * torch.sum(data).unsqueeze(0)) for data in input_data]  # 构造简易数据集
    model = torch.nn.Linear(in_features=16, out_features=1, device='npu')  # 构造简易模型
-   optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3)  # 创建优化器
+   # 使用apex.optimizers.NpuFusedSGD替换torch.optim.SGD
+   optimizer = apex.optimizers.NpuFusedSGD(model.parameters(), lr=1e-3) 
+   # 使用apex.amp.initialize方法包装模型和优化器进行混合精度训练，opt_level='O1'表示使用混合精度训练，combine_grad=True表示开启融合梯度
+   model, optimizer = apex.amp.initialize(models=model, optimizers=optimizer, opt_level='O1', combine_grad=True)
    for i in range(epoch):
        for x, target in dataset:
            optimizer.zero_grad()
            output = model(x)
            loss = loss_function(output, target)
-           loss.backward()
+           # 使用apex.amp.scale_loss方法包装损失值和优化器进行动态损失缩放
+           with apex.amp.scale_loss(loss, optimizer) as scaled_loss:
+               scaled_loss.backward()
            optimizer.step()

特性介绍

Apex-patch已适配特性如下：

• O1/O2自动混合精度模式
• 动态/静态 loss scale
• combine ddp
• combine grad for unscale
• npu fused optimizer: adadelta, adam, adamp, adamw, sgd, lamb, rmsprop, rmsprop_tf
• 动态 loss scale新增可调参数，如：dynamic_init_scale, scale_growth_factor, scale_backoff_factor, scale_window

特性及接口介绍

原生API及参数说明请参考Apex官方文档， Apex-patch新增特性说明如下：

apex.amp

apex.amp.initialize

接口原型

apex.amp.initialize(models, optimizers=None, enabled=True, opt_level="O1", cast_model_type=None, patch_torch_functions=None, keep_batchnorm_fp32=None, master_weights=None, loss_scale=None, cast_model_outputs=None, num_losses=1, verbosity=1, dynamic_init_scale=2.**16, scale_growth_factor=2., scale_backoff_factor=0.5, scale_window=2000, min_loss_scale=None, max_loss_scale=2.**24, combine_grad=None, combine_ddp=None, ddp_replica_count=4, user_cast_preferred=None, check_combined_tensors=None)

接口说明

根据选择的opt_level等配置初始化模型、优化器，也可开启融合梯度优化、融合数据并行优化等。amp.initialize 应在构建完模型和优化器后调用，但应在通过任何 DistributedDataParallel 装饰器装饰模型之前调用。目前，amp.initialize 只应调用一次，尽管它可以处理任意数量的模型和优化器。

新增参数说明

• dynamic_init_scale - 动态loss scale初始值（默认2**16）
• scale_growth_factor - loss scale增长系数（默认2）
• scale_backoff_factor - loss scale回退系数（默认0.5）
• scale_window - loss scale窗口（默认2000）
• combine_grad - 梯度融合开关 （默认None）
• combine_ddp - 融合分布式数据并行 （默认None）
• ddp_replica_count - DDP融合梯度副本数（默认4）
• user_cast_preferred - O1模式下优先选择用户注册的精度模式（默认None）
• check_combined_tensors - tensor融合检查（默认None）

约束条件

启用融合功能（combine_grad/combine_ddp）后，在创建融合张量时会申请融合后张量大小的内存，device内存不足时不建议使用。融合张量内存与原张量共享内存，若更改其一的内存地址，将破坏共享内存机制，可以引起精度异常等问题，使用时须用户自行保证共享内存不被破坏。

示例

model, optim = apex.amp.initialize(model, optim, opt_level="O3", keep_batchnorm_fp32=True, ddp_replica_count=8)

apex.amp.scale_loss

接口原型

apex.amp.scale_loss(loss, optimizers, loss_id=0, model=None, delay_unscale=False, delay_overflow_check=False)

接口说明

使用混合精度时对loss进行scale，避免在低精度模式下梯度溢出。

API及参数说明请参考Apex官方文档，Apex-patch中修改了接口内部实现，以保证在NPU上功能正常，在开启融合功能时使用融合张量进行scale/unscale以提升训练效率。

apex.optimizers

接口说明

融合优化器算法实现上等价于torch中的优化器实现，在梯度更新阶段利用Tensor融合技术，使用融合的梯度和参数进行更新，以提升部分场景下昇腾训练服务器上模型训练的效率。

融合优化器约束条件

启用融合优化器（如apex.optimizers.NpuFusedSGD）后，在创建融合张量时会申请融合后张量大小的内存，device内存不足时不建议使用。融合张量内存与原张量共享内存，若更改其一的内存地址，将破坏共享内存机制，可以引起精度异常等问题，使用时须用户自行保证共享内存不被破坏。

示例

将torch.optim.XXX替换为apex.optimizers.NpuFusedXXX，如：

opt = apex.optimizers.NpuFusedSGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=0.1)

已支持的融合优化器及参数说明

已支持的融合优化器及参数说明如下：

• apex.optimizers.NpuFusedSGD

  class apex.optimizers.NpuFusedSGD(params, lr=required, momentum=MOMENTUM_MIN, dampening=DAMPENING_DEFAULT, weight_decay=WEIGHT_DECAY_MIN, nesterov=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率
  • momentum - 动量（默认值：0.0）
  • dampening - 阻尼系数（默认值：0.0）
  • weight_decay - 权重衰减（默认值：0.0）
  • nesterov - 使用nesterov动量（默认值：False）

• NpuFusedAdam

  class NpuFusedAdam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, amsgrad=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • betas -  用于计算梯度及其平方的运行平均值的系数（默认值：（0.9，0.999））
  • eps - 防止除0，提高数值稳定性 （默认值：1e-8）
  • weight_decay - 权重衰减（默认值：0）
  • amsgrad - 是否使用AMSGrad（默认值：False）

• NpuFusedAdamW

  class NpuFusedAdamW(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=1e-2, amsgrad=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • betas -  用于计算梯度及其平方的运行平均值的系数（默认值：（0.9，0.999））
  • eps - 防止除0，提高数值稳定性 （默认值：1e-8）
  • weight_decay - 权重衰减（默认值：0）
  • amsgrad - 是否使用AMSGrad（默认值：False）

• NpuFusedAdamP

  class NpuFusedAdamP(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0, delta=0.1, wd_ratio=0.1, nesterov=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • betas -  用于计算梯度及其平方的运行平均值的系数（默认值：（0.9，0.999））
  • eps - 分母防除0项，提高数值稳定性 （默认值：1e-8）
  • weight_decay - 权重衰减（默认值：0）
  • delta - 余弦相似度阈值（默认值：0.1）
  • wd_ratio - 权重衰减动态调整速率（默认值：0.1）
  • nesterov - 使用nesterov动量（默认值：False）

• NpuFusedBertAdam

  class NpuFusedBertAdam(params, lr=required, warmup=-1, t_total=-1, schedule='warmup_linear', b1=0.9, b2=0.99, e=1e-6, weight_decay=0.01, max_grad_norm=-1)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • warmup - t_total的warmup比例（默认值：-1，表示不进行warmup）
  • t_total - 学习率调整的步数（默认值：-1，表示固定学习率）
  • schedule - 学习率warmup策略（默认值：'warmup_linear'）
  • b1 - Adams b1（默认值：0.9）
  • b2 - Adams b2（默认值：0.99）
  • e - Adams epsilon（默认值：1e-6）
  • weight_decay - 权重衰减（默认值：0.01）
  • max_grad_norm - 最大梯度正则（默认值：1.0，-1表示不做裁剪）

• NpuFusedAdadelta

  class NpuFusedAdadelta(params, lr=1.0, rho=0.9, eps=1e-6, weight_decay=0)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • rho - 梯度的均方差系数（默认值：0.9）
  • eps - 分母防除0项，提高数值稳定性（默认值：1e-6）
  • weight_decay - 权重衰减（默认值：0）

• Lamb

  class Lamb(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0, adam=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • betas -  用于计算梯度及其平方的运行平均值的系数（默认值：（0.9，0.999））
  • eps - 分母防除0项，提高数值稳定性（默认值：1e-8）
  • weight_decay - 权重衰减（默认值：0）
  • adam - 将strust_ratio设置为1，退化为Adam（默认值：False）

• NpuFusedLamb

  class NpuFusedLamb(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0, adam=False, use_global_grad_norm=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率。（默认值：1e-3）
  • betas -  用于计算梯度及其平方的运行平均值的系数。 （默认值：（0.9，0.999））
  • eps - 分母防除0项，提高数值稳定性（默认值：1e-8）
  • weight_decay - 权重衰减（默认值：0）
  • adam - 将strust_ratio设置为1，退化为Adam（默认值：False）
  • use_global_grad_norm - 使用全局梯度正则（默认值：False）

• NpuFusedRMSprop

  class NpuFusedRMSprop(params, lr=1e-2, alpha=0.99, eps=1e-8, weight_decay=0, momentum=0, centered=False)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率。（默认值：1e-3）
  • alpha - 平滑常量（默认值：0.99）
  • eps - 分母防除0项，提高数值稳定性（默认值：1e-8）
  • weight_decay - 权重衰减（默认值：0）
  • momentum - 动量因子（默认值：0）
  • centered - 计算中心RMSProp（默认值：False）

• NpuFusedRMSpropTF

  class NpuFusedRMSpropTF(params, lr=1e-2, alpha=0.9, eps=1e-10, weight_decay=0, momentum=0., centered=False, decoupled_decay=False, lr_in_momentum=True)

  参数说明：

  • params - 模型参数或模型参数组
  • lr - 学习率（默认值：1e-3）
  • alpha - 平滑常量（默认值：0.9）
  • eps - 分母防除0项，提高数值稳定性（默认值：1e-10）
  • weight_decay - 权重衰减（默认值：0）
  • momentum - 动量因子（默认值：0）
  • centered -  计算中心RMSProp（默认值：False）
  • decoupled_decay - 权重衰减仅作用于参数（默认值：False）
  • lr_in_momentum - 计算动量buffer时使用lr（默认值：True）

Apex-patch代码目录说明

 ├── Apex
     ├──patch
         ├──npu.patch              # Apex-patch对于原生Apex的patch文件，用于原生Apex中混合精度等功能基于昇腾AI处理器的适配
     ├──scripts
         ├──docker/                # 使用Docker构建安装包的脚本与说明文件
         ├──build.sh               # 构建安装包脚本
         ├──gen.sh                 # 使用Apex-patch对官方Apex打patch的脚本
         ├──make_patch.sh          # 生成patch文件脚本
     ├──src
         ├──apex
             ├──contrib/           # 提供Tensor融合的Python API，供融合优化器使用
             ├──optimizers/        # 融合优化器的实现，部分场景下发挥昇腾的算力
         ├──csrc/combine_tensors/  # 提供Tensor融合的C++接口
     ├──tests                      # 测试用例
     ├──LICENSE