CMU 10423 Generative AI：lec18（大模型的分布式训练）

这个文档主要讲解了分布式训练（Distributed Training），特别是如何在多GPU上训练大规模的语言模型。以下是主要内容的概述：

1. 问题背景

• 训练大规模语言模型的主要挑战是内存消耗。

• 训练过程中，内存消耗主要来源于两个方面：

  • 模型权重的存储（通常使用FP16格式）
  • 优化器状态（如动量等）的存储（通常使用FP32格式以保持计算精度）

• 举个例子，对于一个拥有70亿参数的模型，存储其FP16格式的权重需要14GB内存，而存储优化器状态则需要56GB内存。因此，整个模型的内存需求是70GB。对于单个具有80GB内存的GPU（如H100 GPU）来说，无法单独存储和训练这种规模的模型。

2. 解决方案：分布式训练

• 为了解决单个GPU无法容纳大模型的问题，我们需要通过分布式训练（在多GPU上训练）来扩大模型的训练规模。

3. Sharded Optimization（分片优化）

• 核心思想：将模型权重和优化器状态分散存储在多个GPU上。
• 在分片优化中，模型的权重和优化器状态会被分割成多个部分，每个部分分别存储在不同的GPU上。

4. 优化器状态分片（Optimizer Sharding）

• 对于7B（70亿参数）的模型，存储优化器状态需要大量内存。为了解决这一问题，可以将优化器状态分片，分别存储在不同的GPU上。
• 假设有m个GPU，优化器状态可以被分割为m个相等的部分，分别存储在每个GPU上。

5. 完全分片（Fully Sharding）

• 除了优化器状态外，模型的权重也可以被分片存储。通过将模型的权重划分为m部分，分别存储在m个GPU上，可以进一步减小单个GPU的内存负担。

6. 张量并行（Tensor Parallel）

• 与分片不同，张量并行不仅是将权重分割存储，还会将前向传播和反向传播的计算任务分配到不同的GPU上执行。
• 在张量并行中，模型的线性层会被分为行并行和列并行两种方式，通过将矩阵的乘法计算分布到不同的GPU上来加速计算。

7. 流水线并行（Pipeline Parallel）

• 在流水线并行中，模型的不同层会被拆分为一个顺序模块，每个层被分配到不同的GPU上。
• 前向传播和反向传播的计算会依次通过这些层进行，就像流水线一样进行处理。

总结：

分布式训练是训练大规模语言模型的必要手段，尤其当单个GPU无法容纳模型的全部参数时。通过分片优化、张量并行和流水线并行等技术，可以有效地在多个GPU上分担内存和计算任务，保证训练的顺利进行。这些技术使得我们能够训练更大规模的模型，并且显著提高训练速度。