AI模型训练中的多级存储优化

### AI模型训练中的多级存储优化 在AI模型训练中，随着模型规模的不断增大，GPU内存不足的问题日益凸显。为了解决这一问题，研究人员提出了多种方法，下面为你详细介绍。 #### 1. 并行训练技术 目前，有多种张量并行技术可用于AI模型训练，包括1D、2D、2.5D和3D张量并行。此外，流水线并行（PP）考虑了AI模型训练过程的特点，其核心思想是将不同的训练阶段放置在不同的GPU上以实现并行化。不过，PP存在“气泡开销”，而Megatron - LM和PipeDream减少了计算资源的浪费并提高了效率。 #### 2. 异构训练 异构训练是解决GPU内存不足问题的另一种有效方法。微软Zero系列的Zero - Offload和Zero - Infinity分别允许在AI模型训练期间将张量卸载到CPU和非易失性内存快速（NVMe）SSD。TSPLIT将张量划分为更细粒度的单元，并基于此进行内存管理。在异构训练中，需要考虑卸载张量的选择和卸载操作。对于卸载策略，SwapAdviser使用遗传算法进行启发式搜索，AutoTM使用线性整数规划来寻找合适的策略。对于卸载操作，PatrickStar将小张量打包成一个块以提高带宽利用率。 #### 3. 其他训练策略 除了上述两种方法外，还有其他训练大型AI模型的方式： - **梯度累积**：Lin等人使用梯度累积，即将数据分段并依次输入到同一设备进行串行计算，然后将梯度累积并一起更新。 - **激活检查点**：Chen等人使用激活检查点来重新计算部分数据，从而无需保存这些数据。 - **词嵌入缓存数据交换**：对于大语言模型的高维词嵌入，Miao等人设计了服务器和多个客户端之间的词嵌入缓存数据交换，以提高训练过程中分布式系统的通信效率。 #### 4. 多级中间数据卸载策略 提出了一种多级中间数据卸载策略，该策略有效利用CPU内存和HDD来缓解GPU内存压力，包含两个阶段：动态管理策略和HDD卸载。 ##### 4.1 基于预热的动态管理策略 在AI训练过程中，模型数据和非模型数据竞争相同的GPU内存资源，且非模型数据所需的存储空间会动态变化。因此，将GPU内存优化转换为合适的模型数据移动策略，根据非模型数据所需的存储空间动态卸载模型数据。该策略可分为两个步骤： - **动态内存管理**：AI模型训练是一个具有多个epoch的重复计算过程，通常数据集会被分成多个批次。因此，只需利用一次迭代作为预热阶段，就可以了解在固定批次大小下整个训练过程中GPU内存对非模型数据和模型数据的占用情况。在预热阶段，通过以下公式估算非模型数据的GPU内存占用： \[V_{non - model\ data} = V_{max}-\sum Volume(T_{model\ data})\] 其中，\(V_{non - model\ data}\)表示非模型数据的GPU内存占用，\(V_{max}\)表示整个计算过程中的最大内存占用，\(T_{model\ data}\)表示可卸载的模型数据张量，\(Volume(\cdot)\)用于计算模型数据的内存占用。在非预热阶段，根据非模型数据的内存占用动态卸载模型数据。 - **高效卸载策略**：目标是避免类似缓存抖动的中间数据频繁交换。设计了一种类似于Belady的中间数据卸载策略，具体算法如下： ```python Input: hold list: The model data tensor stored in GPU memory; compute list: Tensor usage sequence obtained in the warm - up phase; demand: Capacity required for the non - model data in the next operation. Output: offload queue: The tensor that should be offloaded from GPU memory. 1: for all tensor ∈hold list do 2: if tensor ∈compute list then 3: dict[tensor] = compute list.index(tensor) 4: else 5: dict[tensor] = +∞ 6: end if 7: end for 8: dict.sort(by = values, descending = True) 9: offload queue = [] 10: for all tensor ∈dict.keys() do 11: if avalable ⩾demand then 12: break ```