[发明专利]一种面对AI大模型的分布式张量管理方法

权利要求书

1.一种面对AI大模型的分布式张量管理方法，其特征在于，包括如下步骤：

步骤1、启动N个分布式节点，为每个节点创建单机模型：

步骤2、调用ColoTensor的初始化方法，将单机模型中所有模型参数张量替换为ColoTensor张量；

步骤3、设置模型参数张量的ParallelAction和DistSpec；

步骤4、ColoTensor调用节点间通信重新分配每个模型参数张量的数据，将单机模型自动转化为分布式模型。

步骤5、所有节点执行正向传播；

步骤6、ColoTensor调用节点间通信重新分配每个正向传播要求的输入张量，以匹配该分布式计算需要的输入DistSpec；

步骤7、ColoTensor使用ParallelAction指定的并行方式处理每个节点需要执行的正向传播；

步骤8、ColoTensor调用节点间通信重新分配每个正向传播要求的输出张量，以匹配该分布式计算需要的输出DistSpec；

步骤9、重复步骤7-8直到正向传播完成；

步骤10、所有节点执行反向传播计算；

步骤11、ColoTensor调用节点间通信重新分配每个反向传播计算要求的梯度张量；

步骤12、ColoTensor使用ParallelAction指定的并行方式处理每个节点需要执行的梯度计算；

步骤13、汇总N个节点的数据输出模型结果，结束。

2.根据权利要求1所述的一种面对AI大模型的分布式张量管理方法，其特征在于：所述步骤3中，设置模型参数张量的方式如下：

a、手动设置每个张量的ParallelAction和DistSpec；

b、指定ColoTensor提供的常用并行计算模式一键配置ParallelAction和DistSpec；

c、指定ColoTensor根据模型和计算设备自动推导合适的ParallelAction和DistSpec；

d、混合使用a和b方式。

3.根据权利要求1所述的一种面对AI大模型的分布式张量管理方法，其特征在于：所述自动分配DistSpec的步骤如下：

步骤1、根据用户配置或者自动推导的ParallelAction和DistSpec，ColoTensor自动配置模型参数张量的DistSpec，将其分布在合适的节点计算设备上；

步骤2、正向传播时，ColoTensor根据参与计算的所有模型参数张量的DistSpec，预测需要的输入参数张量DistSpec，并将输入的DistSpec转化；

步骤3、反向传播计算时，ColoTensor根据正向传播时预测的输入输出DistSpec，对梯度做相应计算。

4.根据权利要求1所述的一种面对AI大模型的分布式张量管理方法，其特征在于：所述ColoTensor：分布式张量数据结构，记录张量分布的全局信息与本地数据切片，并封装分布式数据通信，对用户提供像操作单机张量一样的分布式数据操作接口。

5.根据权利要求1所述的一种面对AI大模型的分布式张量管理方法，其特征在于：所述ParallelAction：ColoTensor的属性描述该分布式张量的计算模式，ColoTensor能自动将用户的单机模型转变成并行模型。

6.根据权利要求1所述的一种面对AI大模型的分布式张量管理方法，其特征在于：所述DistSpec：ColoTensor的属性，描述ColoTensor在每个节点具体的分布方式，每个节点重复存储所有张量数据replicate和每个节点分别存储部分张量数据shard两种。

7.根据权利要求6所述的一种面对AI大模型的分布式张量管理方法，其特征在于：所述DistSpecManager：管理ColoTensor在不同DistSpec之间的转换，规定replicate和shard之间以及不同shard之间的转换方式，并封装分布式张量的数据访问。