详解 DeepSpeed Zero 的各个 Stage 状态及日常使用

详解 DeepSpeed Zero 的各个 Stage 状态及日常使用

前文：深度学习模型训练显存占用及DP、MP、PP分布式训练策略

欢迎到访：ckblogs.cn

DeepSpeed 是一个由 Microsoft 开发的深度学习开源优化工具，它横跨模型训练、推理和模型压缩等领域。在本文中，我将只关注与训练相关的优化，而推理和压缩部分将留待今后探讨。

如前文所述，我们可以通过在多个设备上复制整个模型（数据并行 Data Parallelism）或将模型拆分，并将其不同部分存储在不同设备上（模型并行 Model Parallelism / 流水线并行 Pipeline Parallelism）来执行分布式训练。一般来说，DP 比 MP 的计算效率更高；但是，如果模型太大，单个 GPU 设备的可用显存无法容纳，那么只能使用模型并行。

零冗余优化器 Zero Redundancy Optimizer (ZeRO) 是 DeepSpeed 提供的训练优化的核心，它是一套减少分布式模型训练所需内存量的技术。

ZERO

数据并行将模型复制多份至各个 GPU 设备上，但显然这个复制模型的过程将产生较大的显存冗余，为了解决这个问题，有效地降低冗余，可以采用 ZeRO-DP 来取代 DP：

ZeRO-DP 通过以下方式解决这种冗余问题：

1. Partitioning optimizer state （分割优化器状态）
2. Partitioning gradients （划分梯度）
3. Partitioning model parameters （分割模型参数）

这三种冗余的详情在上一篇博文中有具体的介绍

deepSpeed 设置了三种 stage，包括上面这三种解决方案的组合供用户选择：

Stage	描述
Stage 0	不采用任何内存优化方案，也就是普通 DDP
Stage 1	Partitioning optimizer state
Stage 2	Partitioning optimizer state 、Partitioning gradients
Stage 3	Partitioning optimizer state 、Partitioning gradients、Partitioning model parameters

下面，我们聚焦这四种 Stage，分别通过实验方式做了解，一般 deepSpeed 的训练代码结构需要配置 deepspeed.json ，利用 deepspeed 启动命令启动训练，详情请见：Getting Started 。这里为了简要描述，不使用这种方式。

Stage 0

使用 DeepSpeed 执行数据并行训练时，我们不必像使用 PyTorch 的 DistributedDataParallel 那样设置进程组。相反，我们通过调用 deepspeed.initialize，将模型封装在 DeepSpeedEngine 中，由其在内部处理所有分布式训练逻辑。

在这些示例中，我们将使用 HuggingFace 提供的 BLOOM-560M 模型进行预训练。

# deepspeed_stage_0.py
#
# DeepSpeed automatically sets the LOCAL_RANK environment variable
# to the index of the current device.
rank = int(os.getenv("LOCAL_RANK", "0"))

model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m")

deepspeed_config = {
    "train_micro_batch_size_per_gpu": 1
}

model_engine, _, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=deepspeed_config
)

print(f"Device {rank} - ZeRO Stage: {model_engine.zero_optimization_stage()}")

要启动分布式训练作业，我们使用与 deepspeed Python 软件包一起安装的 deepspeed 命令行实用程序：

deepspeed deepspeed_stage_0.py

Device 1 - ZeRO Stage: 0
Device 4 - ZeRO Stage: 0
Device 2 - ZeRO Stage: 0
Device 3 - ZeRO Stage: 0
Device 6 - ZeRO Stage: 0
Device 7 - ZeRO Stage: 0
Device 0 - ZeRO Stage: 0
Device 5 - ZeRO Stage: 0

实际上，deepspeed 会在初始化过程中，创造多个进程，每个进程运行在一张 GPU 上，在 deepspeed_config 中，如果我们不显示的指定 zero-stage，将默认使用 stage 0.

我们可以通过下面的代码开展训练：

# DeepSpeed automatically sets the WORLD_SIZE environment variable
# to the number of devices participating in the training job.
world_size = int(os.getenv("WORLD_SIZE", "1"))

tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-560m")
collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

# helper function to load the wikitext dataset
# implementation can be found here:
# https://github.com/gnovack/distributed-training-and-deepspeed/blob/main/util.py
train_dataset = load_wikitext(tokenizer, collator).select(range(64))

train_dataloader = DataLoader(
    train_dataset,
    batch_size=1,
    shuffle=False,
    sampler=DistributedSampler(train_dataset, num_replicas=world_size)
)

for batch in train_dataloader:
    device = torch.device("cuda", rank)
    
    input_ids = batch['input_ids'].to(device)
    labels = batch['labels'].to(device)

    outputs = model_engine(input_ids, labels=labels)
    model_engine.backward(outputs.loss)
    model_engine.step()

内存占用分析：

根据上篇文章对 Transformer 模型的内存占用公式计算，我们可以估算出存储一个 560M 参数模型的模型参数所需的内存约为 2.24GB。在 32 位模式下使用 Adam 时，优化器状态下每个模型参数占用 8 字节内存，因此 Optimizer States 需要 4.48GB 内存。

假设我们使用 batch size 为 1 和输入序列长度为 512 的模型进行训练，我们可以使用上述公式估算出 Activations 所需的内存为 1.75GB。

因此，在 Stage 0 下，整个模型训练所需要的内存大小如下所示：

我们可以调用 memory_status() 方法来获取每个 GPU 所用的显存：

from deepspeed.runtime.utils import memory_status

for batch in train_dataloader:
    input_ids = batch['input_ids'].to(device)
    labels = batch['labels'].to(device)

    outputs = model_engine(input_ids, labels=labels)
    
    model_engine.backward(outputs.loss)

    model_engine.step()

    if rank == 0:
        memory_status("Memory stats after training step")

在 Stage 0 状态下，得到的显存结果是：

RANK=0 MEMSTATS Memory stats after training step device=cuda:0 current alloc=6.8847GB (delta=0.0000GB max=10.8823GB) current cache=13.4453GB (delta=0.0000GB max=13.4453GB)

这里所占用的最大显存是 max=10.8823GB ，与上文计算得到的 10.71 GB 十分接近。

Stage 1 ：Optimizer State Partitioning

ZeRO Stage 1 将优化器状态（例如，使用 Adam 时，会有第一和第二矩值需要存储）分割到所有设备上，这样每个设备只包含该状态的一部分。

比如下面这两张图，就很好的反应了 Stage 1 的分割结果（上图代表传统的数据并行 Stage 0，下图代表进行了优化器状态分割后的 GPU 显存分配情况）：

实验：

在这里，通过修改 deepspeed config 文件，即可以配置 zero stage 为 1：

deepspeed_config = {
    "train_micro_batch_size_per_gpu": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 5e-5
        }
    },
    "zero_optimization": {
        "stage": 1,
    }
}

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    model_parameters=model.parameters(),
    config=deepspeed_config
)

我们继续使用 memory_status() 方法来获取每个 GPU 所用的显存，可以得到：

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=7.4816GB) current cache=12.7695GB (delta=0.0000GB max=12.7695GB)

现在的 GPU 内存使用峰值为 7.48GB，比我们在使用 Stage 0 训练时测得的 10.88GB 少了 3.4GB，这 3.4 GB 是从哪里少的呢？ 因为我们将原本的优化器状态 4.48 GB，分到了八张 GPU 上，因此，每张 GPU 上的优化器状态所占显存变味了原来的 $\frac{1}{8}$ 也就是 0.56GB：

理论上，使用 Stage 1 得到的现在的 GPU 显存应该为：
$$2.24+2.24+0.56+1.75=6.79GB$$
而事实上的 7.48GB，比这个理论值要大一些，可能是在各个设备之间做相互通信时损失的 GPU 显存。

Stage 2 – Gradient Partitioning

ZeRO Stage 2 更进一步，将优化器状态和梯度都划分到不同的设备上。下图说明了使用 ZeRO Stage 2 进行训练时每个设备的内存分配情况：

deepspeed 的 config 文件应该这样设置：

deepspeed_config = {
    "train_micro_batch_size_per_gpu": 1,
    "optimizer": {
        "type": "Adam",
        "params": {
            "lr": 5e-5
        }
    },
    "zero_optimization": {
        "stage": 2
    }
}

如果在启用 ZeRO Stage 2 的情况下执行训练脚本，我们会发现 GPU 内存的峰值略有增加，达到 8.22GB，而使用阶段 1 进行训练时为 7.48GB：

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=8.2187GB) current cache=12.7715GB (delta=0.0000GB max=12.7715GB)

出现这个问题的原因在于对 reduce_bucket_size 的设置，在每次进行反向传播的过程中，梯度被平均放在了各个设备上，在做梯度下降的过程中，会对梯度打包成一个个 chunk ，再进行梯度下降，当然这样做的原因就在于提高效率，但是当打包的尺寸太大时，也会导致显存占用的增加，控制这个打包尺寸的参数可以通过修改 reduce_bucket_size （默认是 $5\times 10^8$ ） ，我们将其设置为 $5\times 10^6 $ ,可以发现，分配内存下降了不少：

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=3.4843GB (delta=0.0000GB max=6.3750GB) current cache=10.9277GB (delta=0.0000GB max=10.9277GB)

理论上，我们将梯度和优化器显存进行了分割，每个设备所占的显存应该是：
$$2.24+\frac{2.24}{8}+\frac{4.48}{8}+1.75=4.83GB$$
这个理论值和实测值还是有些差距，这一理论估算的前提是，反向传播过程中每次操作产生的梯度都小于 bucket_size 的大小。然而情况并非如此，在反向传播过程中，每个设备都必须先存储这些参数的梯度，然后再通过 Reduce 进行交换，这就导致了一些额外的显存消耗。

其实这里的解释我还是有些发懵，不清楚这里额外的显存消耗到底指的是什么。

除此之外，我们可以通过启用 offload_optimizer 配置减少 GPU 内存使用（需要 "stage": 2）。这个操作的目的是将优化器状态从 GPU 迁移到 CPU 上，但是这个操作也会加大通信时间，而减少 GPU 所占显存：

{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
    }
}

此外，deepspeed==0.4.4 添加了一个新选项 round_robin_gradients，可以通过以下方式启用：

{
    "zero_optimization": {
        "round_robin_gradients": true
    }
}

这是一个用于 CPU offloading 的stage 2优化，通过细粒度梯度分区在 ranks 之间并行复制到 CPU 内存，从而实现了性能的提升。性能优势随着梯度累积步骤（在优化器步骤之间进行更多复制）或 GPU 数量（增加并行性）增加而增加。

Stage 3 – Parameter Partitioning

ZeRO Stage 3在 Stage 2 的基础上对模型参数进行了分区。使用 ZeRO Stage 3 训练 BLOOM 560M 模型时，我们会发现与 Stage 2 相比，每个训练步骤后分配的内存从 3.48GB 降至 1.94GB。但 GPU 的峰值却增加到了 7.48GB。

RANK=0 MEMSTATS Memory stats after training step: device=cuda:0 current alloc=1.9410GB (delta=0.0000GB max=7.4082GB) current cache=10.7129GB (delta=0.0000GB max=10.7129GB)

目前还不清楚 GPU 内存利用率峰值增加的原因，但在 DeepSpeed GitHub 代码库中有一个相关问题的 issue ： DeepSpeed/issues/3734

除此之外，我们也可以通过启用 offload_param 配置减少 GPU 内存使用（需要 "stage": 3）。这个操作的目的是将模型参数从 GPU 迁移到 CPU 上，但是这个操作也会加大通信时间，而减少 GPU 所占显存：

{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
         "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
    }
}

激活值显存占用优化

根据前文对模型占用的显存分析，我们利用下面的代码，分析自己模型训练过程中到底是哪里占用较大的显存：

def train(data_engine, model_engine):
    model_engine.train()
    step = 0
    while step < args.steps_per_epoch:
        before = torch.cuda.memory_allocated()
        print("模型前向传播前使用显存为{}".format(before))
        data = data_engine.get_data()
        loss = model_engine(data, labels=data).loss
        after = torch.cuda.memory_allocated()
        print("模型前向传播后使用显存为{}，差值（中间激活）为{}".format(after, after - before))
        model_engine.backward(loss)
        model_engine.step()
        step += 1
    return


if __name__ == "__main__":
    torch.cuda.empty_cache()
    before = torch.cuda.memory_allocated()
    data_engine = prepare_data()
    model_engine = prepare_model()
    after = torch.cuda.memory_allocated()
    print("建立模型后显存变大{}".format(after - before))

通过运行发现，如果我们的模型本身占用的显存不大，比如选用了较小参数模型，那么绝大部分显存是被前向传播过程中产生的激活值所占用，这是很好理解的，激活内存用于存储神经网络在前向传播过程中计算的中间激活值。激活内存的大小与网络的深度和输入数据大小（batch size）有关。如果选用较大的上下文长度用于训练，比如32K 的长上下文，很容易导致显存爆炸。

而网上的众多资料及 deepspeed 的重点优化步骤都聚焦于模型参数、梯度及优化器，然而现实中的占用情况可能是这样的：

所以才显得优化无济于事。面对 Activation Memory 的优化问题，应当采用 Activation Checkpointing， Activation Checkpointing 技术通过以计算量为代价来节省内存，即不保存整个计算图的所有中间激活值以计算反向传播，而是在反向传播时重新计算这些激活值。它可以应用于模型的任何部分。
具体而言，在前向传递中，函数将以 torch.no_grad() 的方式运行，即不保存中间激活值。相反，在前向传递中，将保存输入元组和函数参数。在反向传递中，检索保存的输入和函数，并再次对函数进行前向传递，现在跟踪中间激活值，然后使用这些激活值计算梯度。

        model = AutoModelForCausalLM.from_pretrained("../Qwen1.5-0.5B/")
        model.gradient_checkpointing_enable()

为何 deepspeed 中的 Gradient Checkpointing 不奏效？

Activation Checkpointing和Gradient Checkpointing是指相同方法的两个不同术语。这确实让人感到困惑，但事实就是这样。

Gradient Checkpointing允许通过牺牲速度来换取GPU内存，HF Transformers 模型对 DeepSpeed 的Activation Checkpointing 一无所知，因此如果尝试在 DeepSpeed 配置文件中启用该功能，什么都不会发生。

因此，有两种方法可以利用这个功能：

1. 如果想使用 HF Transformers 模型，可以使用 model.gradient_checkpointing_enable() 或在 HF Trainer 中使用 --gradient_checkpointing，它会自动为启用这个功能。在这里使用了 torch.utils.checkpoint。
2. 如果编写自己的模型并希望使用DeepSpeed的Activation Checkpointing，可以使用规定的API。还可以使用 HF Transformers 的模型代码，将 torch.utils.checkpoint 替换为 DeepSpeed 的API。后者更灵活，因为它允许将前向激活值卸载到CPU内存，而不是重新计算它们。

如何选择最佳性能的ZeRO Stage和 offloads⚓︎

了解了这些不同stages后，现在需要决定使用哪个stage。通常，可以按照下面的规则进行：

• 速度方面（左边比右边快）

stage 0（DDP） > stage 1 > stage 2 > stage 2 + offload > stage 3 > stage3 + offload

• GPU内存使用方面（右边比左边更节省GPU内存）

stage 0（DDP） < stage 1 < stage 2 < stage 2 + offload < stage 3 < stage 3 + offload

所以，当希望在尽量使用较少数量的GPU的同时获得最快的执行速度时，可以按照以下步骤进行。我们从最快的方法开始，如果遇到GPU内存溢出，然后切换到下一个速度较慢但使用的GPU内存更少的方法。以此类推。

首先，将批量大小设置为1（您始终可以使用梯度累积来获得任何所需的有效批量大小）。

1. 启用 --gradient_checkpointing 1（HF Trainer）或直接 model.gradient_checkpointing_enable() - 如果发生OOM（Out of Memory），则执行以下步骤。

2. 首先尝试 ZeRO stage 2。如果发生OOM，则执行以下步骤。

3. 尝试 ZeRO stage 2 + offload_optimizer - 如果发生OOM，则执行以下步骤。

4. 切换到 ZeRO stage 3 - 如果发生OOM，则执行以下步骤。

5. 启用 offload_param 到 cpu - 如果发生OOM，则执行以下步骤。

6. 启用 offload_optimizer 到 cpu - 如果发生OOM，则执行以下步骤。

7. 如果仍然无法适应批量大小为1，请首先检查各种默认值并尽可能降低它们。例如，如果使用 generate 并且不使用宽搜索束，将其缩小，因为它会占用大量内存。

8. 绝对要使用混合半精度而非fp32 - 在Ampere及更高的GPU上使用bf16，在旧的GPU体系结构上使用fp16。

9. 如果仍然发生OOM，可以添加更多硬件或启用ZeRO-Infinity - 即切换 offload_param 和 offload_optimizer 到 nvme。您需要确保它是非常快的NVMe。作为趣闻，我曾经能够在一个小型GPU上使用BLOOM-176B进行推理，使用了ZeRO-Infinity，尽管速度非常慢。但它奏效了！

当然，也可以按相反的顺序进行这些步骤，从最节省GPU内存的配置开始，然后逐步反向进行，或者尝试进行二分法。

本文章主要内容部分翻译于： https://tinkerd.net/blog/machine-learning/distributed-training/

文章内容融合了自己的思考并调节了部分章节的顺序，请结合原英文内容一同参考。

参考内容

[1] https://tinkerd.net/blog/machine-learning/distributed-training/

[2] DeepSpeed集成

[3] DeepSpeed Docs

[4] deepspeed.ai

[5] Rajbhandari, Samyam, et al. “Zero: Memory optimizations toward training trillion parameter models.” SC20: International Conference for High Performance Computing, Networking, Storage and Analysis. IEEE, 2020.