(19)国家知识产权局 (12)发明 专利申请 (10)申请公布号 (43)申请公布日 (21)申请 号 202210501949.4 (22)申请日 2022.05.10 (71)申请人 北京潞晨科技有限公司 地址 100000 北京市海淀区海淀大街34 号8 层801-802室 (72)发明人 方佳瑞　尤鹏　 (51)Int.Cl. G06F 9/50(2006.01) G06F 17/16(2006.01) G06N 3/04(2006.01) (54)发明名称 一种面向大模型训练的张量存 储管理方法 (57)摘要 本发明公开了一种面向大模型训练的张量 存储管理方法， 涉及计算机技术领域， 具体为一 种面向大模 型训练的张量存储管 理方法， 包括大 模型异构训练的内存管理软件， 所述内存管理软 件包括内存 管理器和内存信息统计器， 所述内存 管理软件让张量在训练过程中动态分布在CPU ‑ GPU的存储空间内， 从而让模型训练突破GPU的内 存墙， 所述内存管理器负责模型数据张量， 所述 内存管理器用于标记所述张量的状态信息， 所述 大模型训练包括预热阶段和正式阶段， 定时采集 系统的CPU或者 GPU的内存使用情况， 并根据每组 的统计时刻， 精确统计非模型数据数值， 通过管 理张量在CP U和GPU的存储方式， 使大模型训练突 破内存墙， 在相同存储硬件配置下， 可 以完成更 大模型的训练。 权利要求书1页 说明书4页 附图1页 CN 114860445 A 2022.08.05 CN 114860445 A 1.一种面向大模型训练 的张量存储管理方法， 其特征在于， 包括方法包括： 大模型异构 训练的内存管理软件， 所述内存管理软件包括内存管理器和内存信息统计器， 所述内存管 理软件让张量在训练过程中动态分布在CPU ‑GPU的存储空间内， 从而 让模型训练突破GPU的 内存墙， 所述内存管理器负责模型数据张量， 所述内存管理器用于标记所述张量的状态信 息， 所述大模型训练包括预 热阶段和正式阶段。 2.根据权利要求1所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 内存管理器包括 查询内存 模组、 转换张量状态和调整张量 位置。 3.根据权利要求2所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 查询内存 模组通过内存管理器获取模型 数据对CPU和GPU的内存占用。 4.根据权利要求2所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 转换张量状态可将张量标记为COMPU TE状态， 且计算之后标记为HOLD状态。 5.根据权利要求2所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 调整张量 位置用于调整张量在设备 上存储位置。 6.根据权利要求1所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 内存信息统计器中数据类型设置有模型数据和非模型数据， 且内存信息统计器用于监测 CPU和GPU中模型 数据和非模型 数据内存使用情况。 7.根据权利要求6所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 模型数据由参数、 梯度和优化器状态组成， 所述非模型数据由算子生成的中间张量和算子 的临时变量组成， 所述模型 数据和非模型 数据相互竞争GPU内存。 8.根据权利要求6所述的一种面向大模型训练 的张量存储管理方法， 其特征在于： 所述 内存信息统计 器通过采样方式在预 热阶段获得非模型对CPU和GPU内存的使用情况。权　利　要　求　书 1/1 页 2 CN 114860445 A 2一种面向大模型训练 的张量存 储管理方法 技术领域 [0001]本发明涉及计算机技 术领域， 具体为 一种面向大模型训练的张量存 储管理方法。 背景技术 [0002]现今以BERT、 GPT为代表的预训练模型(PTM)的出现是自然语言处理(NLP)领域的 里程碑事件。 NLP正在进入预训练时代。 PTM使用堆叠多Transformer结构的神经网络， 在大 量文本上预训练通用语言特征表示， 然后通过微调将学到的知识转移到不同的下游任务。 使用自互联网的海量文本数据， PTM可以捕获自然语 言的细微特征， 并使 下游任务 获得更加 惊艳的表现效果。 因此， A I社区的共识是采用PTM作为特定NLP任务的主干， 而不是在与任务 相关的数据集上从头开始训练模型。 [0003]PTM的力量源泉是它拥有的数以亿计 的参数规模， 这对运行它的计算与内存资源 都提出极高的要求。 由于PTM模型数据无法容纳在单个GPU的内存中， 因此大多数深度学习 训练框架所采用的数据并行技术不适用于PTM。 GPU硬件的有限的存储规模上限， 像一堵墙 限制PTM的可训练规模， 因此从业人员通常称之为 “GPU内存墙 ”现象。 近两年， 业界通过利用 并行训练在多个GPU内存之间分配模型数据， 例如ZeRO ‑DP、 模型并行， 流水线并行尝试使 PTM大小突破内存墙限制。 但是， 这些技术仍严重依赖不断扩 大GPU规模， 而异构训练技术则 成为破局者。 它不仅可以显著 提升单GPU训练模 型的规模， 而且可以和并行训练技术正交使 用。 [0004]异构训练通过在CPU和GPU中容纳模型数据， 并仅在必要时将数据移动到当前设 备， 同时利用GPU内存、 CPU内存(由DRAM或NVMe内存组成)训练大模型。 其他方案如数据并 行、 模型并行、 流水线并行都可以在异构训练基础上进一 步扩展GPU规模。 [0005]目前最佳的异构训练方案DeepSpeed仍存在很大优化空间， 在不考虑非模型数据 的情况下， DeepSpeed在CPU和GPU内存之间静态划分模 型数据， 并且它们的内存布局对于不 同的训练配置是恒定的， 当GPU内存或CPU内存不足以满足其相应的模型数据要求时， 即使 当时其他设备上仍有可用内存， 系统会出现崩溃情况。 发明内容 [0006]本发明的目的在于提供一种面向大模型训练的张量存储管理方法， 以解决上述背 景技术中提出的目前最佳的异构训练方案DeepSpeed仍存在很大优化空间， 在不考虑非模 型数据的情况下， DeepSpeed在CPU和GPU内存之间静态划分模 型数据， 并且它们的内存布局 对于不同的训练配置是恒定的， 当GPU内存或CPU内存不足以满足其相应的模型数据要求 时， 即使当时其 他设备上仍有可用内存， 系统会出现崩溃情况的问题。 [0007]为实现上述目的， 本发明提供如下技术方案： 一种面向大模型训练的张量存储管 理方法， 包括大模型异构训练的内存管理软件， 所述内存管理软件包括内存管理器和内存 信息统计器， 所述内存 管理软件 让张量在训练过程中动态分布在CPU ‑GPU的存储空间内， 从 而让模型训练突破GPU的内存墙， 所述内存管理器负责模 型数据张量， 所述内存 管理器用于说　明　书 1/4 页 3 CN 114860445 A 3