基于MindSpore Serving部署分布式推理

概述
分布式推理是指推理階段采用多卡進行推理，針對超大規模神經網路模型引數個數過多、模型無法完全加載至單卡中進行推理的問題，可利用多卡進行分布式推理，本文介紹部署分布式推理服務的流程，

分布式推理服務的架構如圖所示：

Main行程提供客戶端訪問的介面，管理Distributed Worker并進行任務管理與分發；Distributed Worker行程根據模型配置自動調度Agent完成分布式推理；每一個Agent行程包含一個分布式模型的切片，占用一個device，加載模型執行推理，
上圖展示了rank_size為16，stage_size為2的場景，每個stage包含8個Agent，占用8個device，rank_size表示推理使用的device的個數，stage表示流水線的一段，stage_size表示流水線的段數，Distributed Worker向Agent發送推理請求并從Agent獲取推理結果，Agent之間使用HCCL通信，
當前對分布式模型有以下限制：
1.第一個stage的模型接收相同的輸入資料
2.其他的stage的模型不接收資料，
3.最后一個stage的所有模型都回傳相同的資料，
4.僅支持Ascend 910推理，
下面以一個簡單的分布式網路MatMul為例，演示部署流程，
匯出分布式模型
，需要如下檔案串列：

export_model
├── distributed_inference.py
├── export_model.sh
├── net.py
└── rank_table_8pcs.json

?net.py為MatMul網路定義，
?distributed_inference.py配置分布式相關的引數，
?export_model.sh在當前機器上創建device目錄并且匯出每個device對應的模型檔案，
?rank_table_8pcs.json為配置當前多卡環境的組網資訊的json檔案，
，構造一個包含MatMul、Neg算子的網路
代碼如下：

import numpy as np
from mindspore import Tensor, Parameter, ops
from mindspore.nn import Cell

class Net(Cell):
def init(self, matmul_size, transpose_a=False, transpose_b=False, strategy=None):
super().init()
matmul_np = np.full(matmul_size, 0.5, dtype=np.float32)
self.matmul_weight = Parameter(Tensor(matmul_np))
self.matmul = ops.MatMul(transpose_a=transpose_a, transpose_b=transpose_b)
self.neg = ops.Neg()
if strategy is not None:
self.matmul.shard(strategy)

def construct(self, inputs):
    x = self.matmul(inputs, self.matmul_weight)
    x = self.neg(x)
    return x

使用distributed_inference.py， 配置分布式模型，
代碼如下：

import numpy as np
from net import Net
from mindspore import context, Model, Tensor, export
from mindspore.communication import init

def test_inference():
“”“distributed inference after distributed training”""
context.set_context(mode=context.GRAPH_MODE)
init(backend_name=“hccl”)
context.set_auto_parallel_context(full_batch=True, parallel_mode=“semi_auto_parallel”,
device_num=8, group_ckpt_save_file="./group_config.pb")

predict_data = create_predict_data()
network = Net(matmul_size=(96, 16))
model = Model(network)
model.infer_predict_layout(Tensor(predict_data))
export(model.predict_network, Tensor(predict_data), file_name="matmul", file_format="MINDIR")

def create_predict_data():
“”“user-defined predict data”""
inputs_np = np.random.randn(128, 96).astype(np.float32)
return Tensor(inputs_np)
使用export_model.sh，匯出分布式模型，執行成功后會在上一級目錄創建model目錄，結構如下：

每個device目錄都包含兩個檔案group_config.pb和matmul.mindir，分別表示模型分組組態檔與模型檔案，
部署分布式推理服務
啟動分布式推理服務，需要如下檔案串列：

模型組態檔內容如下：

from mindspore_serving.server import distributed
from mindspore_serving.server import register

model = distributed.declare_servable(rank_size=8, stage_size=1, with_batch_dim=False)

@register.register_method(output_names=[“y”])
def predict(x):
y = register.add_stage(model, x, outputs_count=1)
return y
啟動Serving服務器
呼叫distributed.start_servable方法部署分布式Serving服務器，
代碼如下：

import os
import sys
from mindspore_serving import server
from mindspore_serving.server import distributed

def start():
servable_dir = os.path.dirname(os.path.realpath(sys.argv[0]))
distributed.start_servable(servable_dir, “matmul”,
rank_table_json_file=“rank_table_8pcs.json”,
version_number=1,
distributed_address=“127.0.0.1:6200”,
wait_agents_time_in_seconds=0)

server.start_grpc_server("127.0.0.1:5500")
server.start_restful_server("127.0.0.1:1500")

if name == “main”:
start()
執行推理
通過gRPC訪問推理服務，client需要指定gRPC服務器的網路地址，運行serving_client.py，呼叫matmul分布式模型的predict方法，執行推理，
代碼如下：

import numpy as np
from mindspore_serving.client import Client

def run_matmul():
“”“Run client of distributed matmul”""
client = Client(“localhost:5500”, “matmul”, “predict”)
instance = {“x”: np.ones((128, 96), np.float32)}
result = client.infer(instance)
print(“result:\n”, result)

if name == ‘main’:
run_matmul()
結果如下：