apex:基于昇腾AI处理器的Apex代码补丁项目

Ascend apex adapter

分支26Tags15

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