分布式推理

概述

分布式推理是指推理阶段采用多卡进行推理。如果训练时采用数据并行或者模型参数是合并保存,那么每张卡均持有完整的权重,每张卡推理自身的输入数据,推理方式与单卡推理 一致,只需要注意每卡加载同样的CheckPoint文件进行推理。 本篇教程主要介绍在多卡训练过程中,每张卡上保存模型的切片,在推理阶段采用多卡形式,按照推理策略重新加载模型进行推理的过程。针对超大规模神经网络模型的参数个数过多,模型无法完全加载至单卡中进行推理的问题,可利用多卡进行分布式推理。

当模型非常大,本教程中使用load_distributed_checkpoint 接口主机内存不足情况下,可以参考分布式弹性训练与推理 章节,采用每张卡加载自身对应的切片Checkpoint的方式。

操作实践

样例代码说明

分布式推理样例代码:distributed_inference

操作流程

  1. 执行训练,生成CheckPoint文件和模型参数切分策略文件。

    • 分布式训练教程和样例代码可参考分布式并行训练基础样例(Ascend)

    • 在分布式推理场景中,训练阶段的CheckpointConfig接口的integrated_save参数需设定为False,表示每卡仅保存模型切片而不是全量模型。

    • set_auto_parallel_context接口的parallel_mode参数需设定为auto_parallel或者semi_auto_parallel,并行模式为自动并行或者半自动并行。

    • 此外还需指定strategy_ckpt_save_file参数,即生成的策略文件的地址。

    • 若采用流水线分布式推理,则训练也必须采用流水线并行训练,并且流水线并行训练和推理所用的device_num以及pipeline_stages必须相同。流水线并行推理时,micro_batch为1,不需要调用PipelineCell。参考流水线并行训练教程。

  2. 设置context,根据推理数据推导出推理策略。

    set_auto_parallel_context(full_batch=True, parallel_mode='semi_auto_parallel', strategy_ckpt_load_file='./train_strategy.ckpt')
    network = Net()
    model = Model(network)
    predict_data = create_predict_data()
    predict_strategy = model.infer_predict_layout(predict_data)
    

    其中,

    • full_batch:是否全量导入数据集,为True时表明全量导入,每卡的数据相同,该场景中必须设置为True

    • parallel_mode:并行模式,该场景中必须设置为自动并行或者半自动并行模式。

    • strategy_ckpt_load_file:训练阶段生成的策略文件的文件地址,分布式推理场景中该参数必须设置。

    • create_predict_data:用户需自定义的接口,返回推理数据。与训练阶段不同的是,分布式推理场景中返回类型必须为Tensor

    • infer_predict_layout:根据推理数据生成推理策略。

  3. 导入CheckPoint文件,根据推理策略加载相应的模型切片至每张卡中。

    ckpt_file_list = create_ckpt_file_list()
    load_distributed_checkpoint(network, ckpt_file_list, predict_strategy)
    

    其中,

    • create_ckpt_file_list:用户需自定义的接口,返回按rank id排序的CheckPoint文件名列表。

    • load_distributed_checkpoint:对模型切片进行合并,再根据推理策略进行切分,加载至网络中。

    对于流水线并行推理,每个stage只需要加载本stage的CheckPoint文件。

    load_distributed_checkpoint接口支持predict_strategy为None,此时为单卡推理,其过程与分布式推理有所不同,详细用法请参考链接

  4. 进行推理,得到推理结果。

    model.predict(predict_data)
    

分布式场景导出MindIR文件

在超大规模神经网络模型的场景中,针对因为参数量过大,导致模型无法进行单卡推理的问题,可以采用分布式推理方案。此时在运行推理任务前,需要导出多个MindIR文件。具体方法如下:

首先,需要准备CheckPoint文件和训练策略文件。

CheckPoint文件在训练过程中产生。CheckPoint具体用法可参考: CheckPoint用法

训练策略文件,需要在训练时通过设置context生成,context配置项如下:

set_auto_parallel_context(strategy_ckpt_save_file='train_strategy.ckpt')

这样在训练后,就会在设置的目录下产生名为train_strategy.ckpt的训练策略文件。

由于导出MindIR文件前,一般需要加载CheckPoint文件,而加载分布式训练的CheckPoint文件,需要结合训练策略和推理策略,所以还需生成推理策略文件。 产生推理策略的代码如下:

predict_strategy = model.infer_predict_layout(predict_data)

然后,使用加载分布式CheckPoint的方法,把之前训练好的参数,加载到网络中。 代码如下:

load_distributed_checkpoint(model, ckpt_file_list, predict_strategy)

最后,就可以导出在分布式推理场景下的MindIR文件。

核心代码如下:

# 在context中配置训练过程中产生的策略文件
set_auto_parallel_context(strategy_ckpt_load_file='train_strategy.ckpt')
# 定义网络结构
network = Net()
model = Model(network)
# 得到推理策略文件
predict_strategy = model.infer_predict_layout(predict_data)
# 创建checkpoint list
ckpt_file_list = create_ckpt_file_list()
# 加载分布式参数
load_distributed_checkpoint(model, ckpt_file_list, predict_strategy)
# 导出分布式MindIR文件
export(net, Tensor(input), file_name='net', file_format='MINDIR')

多卡训练、单卡推理的情况,导出MindIR的用法与单机相同。

分布式场景导出MindIR文件样例代码:distributed_export