模型加载
概述
分布式下的模型加载主要是指分布式推理,即推理阶段采用多卡进行推理。如果训练时采用数据并行或者模型参数是合并保存,那么每张卡均持有完整的权重,每张卡推理自身的输入数据,推理方式与单卡推理一致,只需要注意每卡加载同样的CheckPoint文件进行推理。 本篇教程主要介绍在多卡训练过程中,每张卡上保存模型的切片,在推理阶段采用多卡形式,按照推理策略重新加载模型进行推理的过程。针对超大规模神经网络模型的参数个数过多,模型无法完全加载至单卡中进行推理的问题,可利用多卡进行分布式推理。
当模型非常大,本教程中使用load_distributed_checkpoint接口主机内存不足情况下,可以参考模型转换 章节,先对网络执行编译,然后执行分布式Checkpoint转换,就可以让每张卡加载自身对应的切片Checkpoint。
若采用流水线分布式推理,则训练也必须采用流水线并行训练,并且流水线并行训练和推理所用的
device_num
以及pipeline_stages
必须相同。流水线并行推理时,micro_batch
为1,不需要调用PipelineCell
,每个stage
只需要加载本stage
的Checkpoint文件。参考流水线并行训练教程。
相关接口:
mindspore.set_auto_parallel_context(full_batch=True, parallel_mode=ParallelMode.SEMI_AUTO_PARALLEL)
:设置并行配置,其中full_batch
表示是否全量导入数据集,为True
时表明全量导入,每卡的数据相同,该场景中必须设置为True
。parallel_mode
为并行模式,该场景中必须设置为自动并行或者半自动并行模式。mindspore.set_auto_parallel_context(strategy_ckpt_config=strategy_ckpt_dict)
:用于设置并行策略文件的配置。strategy_ckpt_dict
是用于设置并行策略文件的配置,是字典类型。strategy_ckpt_dict = {“load_file”: “./stra0.ckpt”, “save_file”: “./stra1.ckpt”, “only_trainable_params”: False},其中:load_file(str)
:加载并行切分策略的路径,训练阶段生成的策略文件的文件地址,分布式推理场景中该参数必须设置。默认值:""
。save_file(str)
:保存并行切分策略的路径。默认值:""
。only_trainable_params(bool)
:仅保存/加载可训练参数的策略信息。默认值:True
。
model.infer_predict_layout(predict_data)
:根据推理数据生成推理策略。model.infer_train_layout(data_set)
:根据训练数据生成训练策略。mindspore.load_distributed_checkpoint(network, checkpoint_filenames, predict_strategy)
:该接口对模型切片进行合并,再根据推理策略进行切分,加载至网络中。其中network
为网络结构,checkpoint_filenames
是Checkpoint文件的名称构成的list,按rank id顺序排列。predict_strategy
是model.infer_predict_layout
导出的推理策略。
操作实践
下面以单机8卡为例,进行分布式推理的操作说明。
样例代码说明
下载完整的样例代码:model_saving_loading。
目录结构如下:
└─ sample_code
├─ model_saving_loading
├── test_loading.py
├── run_loading.sh
...
...
其中,test_loading.py
是定义网络结构和推理的脚本。run_loading.sh
是执行脚本。
用户首先需要按照模型保存教程执行8卡分布式训练,训练结束后将会在当前路径生成Checkpoint文件目录以及切分策略文件:
src_checkpoints/
src_strategy.ckpt
配置分布式环境
通过context接口指定运行模式、运行设备、运行卡号等,与单卡脚本不同,并行脚本还需指定并行模式parallel_mode
为半自动并行模式,通过strategy_ckpt_config
配置需要加载的分布式策略文件路径,并通过init初始化HCCL或NCCL通信。device_target
会自动指定为MindSpore包对应的后端硬件设备。
import mindspore as ms
from mindspore.communication import init
ms.set_context(mode=ms.GRAPH_MODE)
ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.SEMI_AUTO_PARALLEL)
ms.set_auto_parallel_context(strategy_ckpt_config={"load_file": "./src_strategy.ckpt"})
init()
ms.set_seed(1)
网络定义
推理网络定义需要与训练网络相同:
from mindspore import nn, ops
from mindspore.common.initializer import initializer
class Dense(nn.Cell):
def __init__(self, in_channels, out_channels):
super().__init__()
self.weight = ms.Parameter(initializer("normal", [in_channels, out_channels], ms.float32))
self.bias = ms.Parameter(initializer("normal", [out_channels], ms.float32))
self.matmul = ops.MatMul()
self.add = ops.Add()
def construct(self, x):
x = self.matmul(x, self.weight)
x = self.add(x, self.bias)
return x
class Network(nn.Cell):
def __init__(self):
super().__init__()
self.flatten = ops.Flatten()
self.layer1 = Dense(28*28, 512)
self.relu1 = ops.ReLU()
self.layer2 = Dense(512, 512)
self.relu2 = ops.ReLU()
self.layer3 = Dense(512, 10)
def construct(self, x):
x = self.flatten(x)
x = self.layer1(x)
x = self.relu1(x)
x = self.layer2(x)
x = self.relu2(x)
logits = self.layer3(x)
return logits
net = Network()
net.layer1.matmul.shard(((2, 1), (1, 2)))
net.layer3.matmul.shard(((2, 2), (2, 1)))
推理
分布式推理与单机模式推理的区别在于,分布式推理需要调用model.infer_predict_layout
接口和ms.load_distributed_checkpoint
接口加载分布式参数,且输入类型必需为Tensor。代码如下:
import numpy as np
import mindspore as ms
# 分布式推理场景输入类型必须为Tensor
predict_data = ms.Tensor(np.random.randn(1, 28, 28).astype(np.float32))
model = ms.Model(net)
# 得到推理策略文件
predict_layout = model.infer_predict_layout(predict_data)
# 创建checkpoint list
ckpt_file_list = ["./src_checkpoints/rank_{}/checkpoint-10_1875.ckpt".format(i) for i in range(0, get_group_size())]
# 加载分布式参数
ms.load_distributed_checkpoint(net, ckpt_file_list, predict_layout)
predict_result = model.predict(predict_data)
print(predict_result)
分布式场景导出MindIR文件
在超大规模神经网络模型的场景中,针对因为参数量过大,导致模型无法进行单卡推理的问题,可以采用分布式推理方案。此时在运行推理任务前,需要导出多个MindIR文件。核心代码如下:
import mindspore as ms
from mindspore.communication import get_rank
# 导出分布式MindIR文件
ms.export(net, predict_data, file_name='./mindir/net_rank_' + str(get_rank()), file_format='MINDIR')
多卡训练、单卡推理的情况,导出MindIR的用法与单机相同。
运行单机八卡脚本
接下来通过命令调用对应的脚本,以mpirun
启动方式,8卡的分布式推理脚本为例,进行分布式推理:
bash run_loading.sh
推理结果保存在log_output/1/rank.*/stdout
中,示例如下:
[[ 0.0504479 -0.94413316 0.84689146 -0.28818333 0.66444737 1.0564338
-0.04191194 0.25590336 -0.69010115 -0.6532427 ]]
其中MindIR文件保存在mindir
目录中,目录结构如下:
├─ mindir
| ├─ net_rank_0.mindir
| ├─ net_rank_1.mindir
| ├─ net_rank_2.mindir
| ├─ net_rank_3.mindir
| ├─ net_rank_4.mindir
| ├─ net_rank_5.mindir
| ├─ net_rank_6.mindir
| └─ net_rank_7.mindir
...