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

说明书

本发明公开了一种面对AI大模型的分布式张量管理方法，涉及张量管理技术领域，具体为一种面对AI大模型的分布式张量管理方法，包括通过数据并行和模型并行的两种并行训练和推理方式，一方面可通过切分模型的输入张量，使得各个机器上有相同的模型但是处理不同的输入，从而减少总体的训练时间，另一方面可通过切分模型的参数张量，使得模型均匀地分布在各个机器上，从而减少每个机器的内存开销，均摊所有计算操作，将分布式张量底层的数据移动自动封装和处理，对用户提供张量粒度的分布式方案，并通过对分布式训练推理的核心数据结构张量，做有效抽象和管理，从而提供用户从单机到分布式模型训练推理的便捷转化。

技术领域

本发明涉及张量管理技术领域，具体为一种面对AI大模型的分布式张量管理方法。

背景技术

新的深度学习模型不断对各种任务的效果做出了突破。这些模型依赖于深而宽的网络结构来“记忆”某些表征，往往把模型的网络增多增宽，增加参数量能进一步提升模型表现，因此近年来模型的模型参数也越来越多。比如Bert论文中提到的bert-base有1亿参数和bert-large有3亿参数；最近的GPT-3和PaLM的large更是高达1750亿参数和5400亿参数。

大规模模型为分布式训练和推理带来了新的挑战：过去单卡就能放下的小规模模型仅用单机多卡训练就能达到优秀的性能和可扩展性；而如今单个模型甚至算子的参数量就超过了单卡内存，需要用更复杂的并行技术来将参数分布到各个节点多张卡上，使分布式训练推理能支持更大规模的模型。

张量作为神经网络的基本数据结构，在分布式时代同样需要被重新思考和设计，分布式的数据并行和模型并行都是对非模型参数张量和模型参数张量做了分布式的切分存储，为用户提供张量粒度的分布式管理是非常重要的，可以说AI大模型的分布式处理底层其实就是对系统中各类张量做了分布式处理。但是目前领域内往往着眼于思考框架和模型层面的方案，对张量这个基本数据结构的关注不够。

现有的AI框架如PyTorch、Tensorflow设计服务于单机，张量管理方案是本地单机的张量数据结构，缺少对分布式训练推理的支持，使用这些单机张量管理方案来管理分布式张量需要很强的专业知识；而已有的分布式训练推理技术如Deepspeed等从模型层面做专业改造，为用户提供改造后的分布式模型，没有提供对分布式张量的管理，因为技术关注的粒度较粗，用户使用这些技术定制自己的模型难度较大。

发明内容

本发明的目的在于提供一种面对AI大模型的分布式张量管理方法，以解决上述背景技术中提出的现有的AI框架如PyTorch、Tensorflow设计服务于单机，张量管理方案是本地单机的张量数据结构，缺少对分布式训练推理的支持，使用这些单机张量管理方案来管理分布式张量需要很强的专业知识；而已有的分布式训练推理技术如Deepspeed等从模型层面做专业改造，为用户提供改造后的分布式模型，没有提供对分布式张量的管理，因为技术关注的粒度较粗，用户使用这些技术定制自己的模型难度较大的问题。

为实现上述目的，本发明提供如下技术方案：一种面对AI大模型的分布式张量管理方法，包括如下步骤：

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

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

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

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

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

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