DeepSpeed零冗余优化器Zero Redundancy Optimizer

阿里云国内75折回扣微信号：monov8

阿里云国际，腾讯云国际，低至75折。AWS 93折免费开户实名账号代冲值优惠多多微信号：monov8 飞机：@monov6

零冗余优化器

内容

如果您还没有这样做我们建议您在逐步完成本教程之前先阅读有关入门和Megatron-LM GPT-2 的DeepSpeed 教程。

在本教程中我们将 ZeRO 优化器应用于Megatron-LM GPT-2模型。ZeRO 是一组强大的内存优化技术可以有效训练具有数万亿参数的大型模型例如GPT-2和Turing-NLG 17B。与训练大型模型的替代模型并行方法相比ZeRO 的一个关键吸引力在于不需要修改模型代码。正如本教程将演示的那样在 DeepSpeed 模型中使用 ZeRO 既快速又简单因为您只需更改 DeepSpeed 配置 JSON 中的一些配置。无需更改代码。

零概述固定链接

ZeRO 利用数据并行的聚合计算和内存资源来减少用于模型训练的每个设备 (GPU) 的内存和计算需求。ZeRO 通过在分布式训练硬件中的可用设备GPU 和 CPU之间划分各种模型训练状态权重、梯度和优化器状态来减少每个 GPU 的内存消耗。具体而言ZeRO 正在作为优化的增量阶段实施其中早期阶段的优化在后期阶段可用。要深入了解 ZeRO请参阅我们的论文。

阶段 1优化器状态例如对于Adam 优化器、32 位权重以及第一和第二矩估计在进程之间进行分区以便每个进程仅更新其分区。
阶段 2用于更新模型权重的减少的 32 位梯度也被分区这样每个进程只保留与其优化器状态部分对应的梯度。
第 3 阶段16 位模型参数跨进程进行分区。ZeRO-3 将在前向和后向传递过程中自动收集和划分它们。

此外ZeRO-3 包括无限卸载引擎以形成 ZeRO-Infinity论文它可以卸载到 CPU 和 NVMe 内存以节省大量内存。

培训环境固定链接

我们在本练习中使用 DeepSpeed Megatron-LM GPT-2 代码。您可以逐步完成 Megatron-LM教程以熟悉代码。我们将在具有 32GB RAM 的NVIDIA Tesla V100-SXM3 Tensor Core GPU上训练本教程中的模型。

启用零优化固定链接

要为 DeepSpeed 模型启用 ZeRO 优化我们只需将zero_optimization键添加到 DeepSpeed JSON 配置。此处提供了zero_optimization键的配置旋钮的完整描述。

训练 1.5B 参数 GPT-2 模型固定链接

我们展示了 ZeRO 第 1 阶段的优势展示了它可以在八个 V100 GPU 上对 15 亿参数的 GPT-2 模型进行数据并行训练。我们将训练配置为使用每个设备 1 的批处理大小以确保内存消耗主要是由于模型参数和优化器状态。我们通过对 deepspeed 启动脚本应用以下修改来创建此训练场景

<span style="background-color:#263238"><span style="color:#eeffff"><code>       <span style="color:#89ddff">--model-parallel-size</span> 1 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--num-layers</span> 48 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--hidden-size</span> 1600 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--num-attention-heads</span> 16 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--batch-size</span> 1 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--deepspeed_config</span> ds_zero_stage_1.config <span style="color:#f78c6c">\</span>
</code></span></span>

在没有 ZeRO 的情况下训练此模型失败并出现内存不足 (OOM) 错误如下所示

该模型不适合 GPU 内存的一个关键原因是模型的 Adam 优化器状态消耗 18GB32GB RAM 的很大一部分。通过使用 ZeRO 第 1 阶段将优化器状态划分为八个数据并行列每个设备的内存消耗可以减少到 2.25GB从而使模型可训练。要启用 ZeRO 阶段 1我们只需更新 DeepSpeed JSON 配置文件如下所示

<span style="background-color:#263238"><span style="color:#eeffff"><code><span style="color:#eeffff">{</span>
    <span style="color:#eeffff">"zero_optimization"</span><span style="color:#eeffff">:</span> <span style="color:#eeffff">{</span>
        <span style="color:#eeffff">"stage"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"reduce_bucket_size"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">5e8</span>
    <span style="color:#eeffff">}</span>
<span style="color:#eeffff">}</span>
</code></span></span>

如上所示我们在zero_optimization键中设置了两个字段。具体来说我们将stage字段设置为 1并将梯度缩减的可选reduce_bucket_size设置为 500M。启用 ZeRO stage 1 后该模型现在可以在 8 个 GPU 上顺利训练而不会耗尽内存。下面我们提供一些模型训练的截图

从nvidia-smi上面的屏幕截图我们可以看到只有 GPU 6-7 用于训练模型。使用 ZeRO stage 1我们可以通过提高数据并行度来进一步降低每台设备的内存消耗。这些内存节省可用于增加模型大小和/或批量大小。相比之下单靠数据并行性无法实现此类优势。

训练 10B 参数 GPT-2 模型固定链接

ZeRO 第 2 阶段优化进一步增加了可以使用数据并行性进行训练的模型的大小。我们通过使用 32 个 V100 GPU 训练具有 10B 个参数的模型来展示这一点。

首先我们需要配置一个启用了激活检查点的 10B 参数模型。这可以通过将以下 GPT-2 模型配置更改应用于 DeepSpeed 启动脚本来完成。

<span style="background-color:#263238"><span style="color:#eeffff"><code>       <span style="color:#89ddff">--model-parallel-size</span> 1 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--num-layers</span> 50 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--hidden-size</span> 4096 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--num-attention-heads</span> 32 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--batch-size</span> 1 <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--deepspeed_config</span> ds_zero_stage_2.config <span style="color:#f78c6c">\</span>
       <span style="color:#89ddff">--checkpoint-activations</span>
</code></span></span>

接下来我们需要更新 DeepSpeed JSON 配置如下所示以启用 ZeRO 阶段 2 优化

<span style="background-color:#263238"><span style="color:#eeffff"><code><span style="color:#eeffff">{</span>
    <span style="color:#eeffff">"zero_optimization"</span><span style="color:#eeffff">:</span> <span style="color:#eeffff">{</span>
        <span style="color:#eeffff">"stage"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">2</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"contiguous_gradients"</span><span style="color:#eeffff">:</span> <span style="color:#c792ea">true</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"overlap_comm"</span><span style="color:#eeffff">:</span> <span style="color:#c792ea">true</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"reduce_scatter"</span><span style="color:#eeffff">:</span> <span style="color:#c792ea">true</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"reduce_bucket_size"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">5e8</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"allgather_bucket_size"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">5e8</span>
    <span style="color:#eeffff">}</span>
<span style="color:#eeffff">}</span>
</code></span></span>

在上述更改中我们将stage字段设置为 2并配置了 ZeRO 阶段 2 中可用的其他优化旋钮。例如我们启用了contiguous_gradients减少反向传递期间的内存碎片。此处提供了这些优化旋钮的完整描述。通过这些更改我们现在可以启动训练运行。

这是训练日志的屏幕截图

这是 nvidia-smi 的屏幕截图显示了训练期间的 GPU 活动

使用 ZeRO-Infinity 训练万亿级模型

ZeRO-3 是 ZeRO 的第三个阶段它对整个模型状态即权重、梯度和优化器状态进行分区以根据数据并行度线性扩展内存节省。可以在 JSON 配置中启用 ZeRO-3。此处提供了这些配置的完整描述。

使用 ZeRO-Infinity 卸载到 CPU 和 NVMe

ZeRO-Infinity 使用 DeepSpeed 的无限卸载引擎将整个模型状态卸载到 CPU 或 NVMe 内存从而允许更大的模型尺寸。可以在 DeepSpeed 配置中启用卸载

<span style="background-color:#263238"><span style="color:#eeffff"><code><span style="color:#eeffff">{</span>
    <span style="color:#eeffff">"zero_optimization"</span><span style="color:#eeffff">:</span> <span style="color:#eeffff">{</span>
        <span style="color:#eeffff">"stage"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">3</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"contiguous_gradients"</span><span style="color:#eeffff">:</span> <span style="color:#c792ea">true</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"stage3_max_live_parameters"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1e9</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"stage3_max_reuse_distance"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1e9</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"stage3_prefetch_bucket_size"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1e7</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"stage3_param_persistence_threshold"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1e5</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"reduce_bucket_size"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1e7</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"sub_group_size"</span><span style="color:#eeffff">:</span> <span style="color:#f78c6c">1e9</span><span style="color:#eeffff">,</span>
        <span style="color:#eeffff">"offload_optimizer"</span><span style="color:#eeffff">:</span> <span style="color:#eeffff">{</span>
            <span style="color:#eeffff">"device"</span><span style="color:#eeffff">:</span> <span style="color:#c3e88d">"cpu"</span>
         <span style="color:#eeffff">},</span>
        <span style="color:#eeffff">"offload_param"</span><span style="color:#eeffff">:</span> <span style="color:#eeffff">{</span>
            <span style="color:#eeffff">"device"</span><span style="color:#eeffff">:</span> <span style="color:#c3e88d">"cpu"</span>
       <span style="color:#eeffff">}</span>
   <span style="color:#eeffff">}</span>
<span style="color:#eeffff">}</span>
</code></span></span>

ZeRO-Infinity 与 ZeRO-Offload DeepSpeed 首先包括 ZeRO-Offload 的卸载功能ZeRO-Offload 是一种用于将优化器和梯度状态卸载到 ZeRO-2 中的 CPU 内存的系统。ZeRO-Infinity 是 ZeRO-3 可访问的下一代卸载功能。ZeRO-Infinity 能够比 ZeRO-Offload 卸载更多的数据并且具有更有效的带宽利用率和计算与通信的重叠。

分配 Massive Megatron-LM 模型固定链接

我们对模型初始化进行了两项进一步更改以支持超过本地系统内存但不超过总系统内存的模型。

以内存可扩展的方式分配模型。模型参数将被分配并立即在数据并行组中进行分区。如果remote_device是"cpu"或"nvme"模型也将分配在 CPU/NVMe 内存中而不是 GPU 内存中。有关详细信息请参阅完整的ZeRO-3 Init 文档。

<span style="background-color:#263238"><span style="color:#eeffff"><code> <span style="color:#c792ea">with</span> <span style="color:#eeffff">deepspeed</span><span style="color:#eeffff">.</span><span style="color:#eeffff">zero</span><span style="color:#eeffff">.</span><span style="color:#eeffff">Init</span><span style="color:#eeffff">(</span><span style="color:#eeffff">data_parallel_group</span><span style="color:#89ddff">=</span><span style="color:#eeffff">mpu</span><span style="color:#eeffff">.</span><span style="color:#eeffff">get_data_parallel_group</span><span style="color:#eeffff">(),</span>
                          <span style="color:#eeffff">remote_device</span><span style="color:#89ddff">=</span><span style="color:#eeffff">get_args</span><span style="color:#eeffff">().</span><span style="color:#eeffff">remote_device</span><span style="color:#eeffff">,</span>
                          <span style="color:#eeffff">enabled</span><span style="color:#89ddff">=</span><span style="color:#eeffff">get_args</span><span style="color:#eeffff">().</span><span style="color:#eeffff">zero_stage</span><span style="color:#89ddff">==</span><span style="color:#f78c6c">3</span><span style="color:#eeffff">):</span>
     <span style="color:#eeffff">model</span> <span style="color:#89ddff">=</span> <span style="color:#eeffff">GPT2Model</span><span style="color:#eeffff">(</span><span style="color:#eeffff">num_tokentypes</span><span style="color:#89ddff">=</span><span style="color:#f78c6c">0</span><span style="color:#eeffff">,</span> <span style="color:#eeffff">parallel_output</span><span style="color:#89ddff">=</span><span style="color:#eeffff">True</span><span style="color:#eeffff">)</span>
</code></span></span>

收集用于初始化的嵌入权重。DeepSpeed 将在其构造函数期间自动收集模块的参数并用于前向和后向传递。但是额外的访问必须与 DeepSpeed 协调以确保收集参数数据并随后对其进行分区。如果修改了张量modifier_rank则还应使用该参数来确保所有等级对数据的看法一致。有关详细信息请参阅完整的GatheredParameters 文档。

<span style="background-color:#263238"><span style="color:#eeffff"><code> <span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">position_embeddings</span> <span style="color:#89ddff">=</span> <span style="color:#eeffff">torch</span><span style="color:#eeffff">.</span><span style="color:#eeffff">nn</span><span style="color:#eeffff">.</span><span style="color:#eeffff">Embedding</span><span style="color:#eeffff">(...)</span>
 <span style="color:#c792ea">with</span> <span style="color:#eeffff">deepspeed</span><span style="color:#eeffff">.</span><span style="color:#eeffff">zero</span><span style="color:#eeffff">.</span><span style="color:#eeffff">GatheredParameters</span><span style="color:#eeffff">(</span><span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">position_embeddings</span><span style="color:#eeffff">.</span><span style="color:#eeffff">weight</span><span style="color:#eeffff">,</span>
                                        <span style="color:#eeffff">modifier_rank</span><span style="color:#89ddff">=</span><span style="color:#f78c6c">0</span><span style="color:#eeffff">):</span>
     <span style="color:#b2ccd6"># Initialize the position embeddings.
</span>     <span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">init_method</span><span style="color:#eeffff">(</span><span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">position_embeddings</span><span style="color:#eeffff">.</span><span style="color:#eeffff">weight</span><span style="color:#eeffff">)</span>

 <span style="color:#eeffff">...</span>

 <span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">tokentype_embeddings</span> <span style="color:#89ddff">=</span> <span style="color:#eeffff">torch</span><span style="color:#eeffff">.</span><span style="color:#eeffff">nn</span><span style="color:#eeffff">.</span><span style="color:#eeffff">Embedding</span><span style="color:#eeffff">(...)</span>
 <span style="color:#c792ea">with</span> <span style="color:#eeffff">deepspeed</span><span style="color:#eeffff">.</span><span style="color:#eeffff">zero</span><span style="color:#eeffff">.</span><span style="color:#eeffff">GatheredParameters</span><span style="color:#eeffff">(</span><span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">tokentype_embeddings</span><span style="color:#eeffff">.</span><span style="color:#eeffff">weight</span><span style="color:#eeffff">,</span>
                                     <span style="color:#eeffff">modifier_rank</span><span style="color:#89ddff">=</span><span style="color:#f78c6c">0</span><span style="color:#eeffff">):</span>
     <span style="color:#b2ccd6"># Initialize the token-type embeddings.
</span>     <span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">init_method</span><span style="color:#eeffff">(</span><span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">tokentype_embeddings</span><span style="color:#eeffff">.</span><span style="color:#eeffff">weight</span><span style="color:#eeffff">)</span>
</code></span></span>

以内存为中心的平铺固定链接

ZeRO-Infinity 包括层的替代品Linear可进一步减少内存。我们可选择平铺模型平行线性层在每个 Transformer 层中找到。注意在构建层时指定相应的基类可以将模型并行和平铺结合起来。该deepspeed.zero.TiledLinear模块利用 ZeRO-3 的数据获取和释放模式通过将大型运算符分解为可按顺序执行的较小块来减少工作内存需求。

我们包括对 Megatron-LM 的ParallelMLP中一个示例的更改。另外三个模型并行层以transformer.py类似方式进行。

Megatron-LM 的模型并行层有一种特殊的形式其中bias层的加法被延迟而不是从forward()后面的运算符中返回以与后面的运算符融合。DeepSpeed的deepspeed.zero.TiledLinearReturnBias子类TiledLinear也简单的将返回的bias参数转发不累加。

<span style="background-color:#263238"><span style="color:#eeffff"><code><span style="color:#eeffff">@@ -1,6 +1,9 @@</span>
<span style="color:#f07178">-self.dense_h_to_4h = mpu.ColumnParallelLinear(
</span><span style="color:#c3e88d">+self.dense_h_to_4h = deepspeed.zero.TiledLinearReturnBias(
</span>     args.hidden_size,
     4 * args.hidden_size,
<span style="color:#c3e88d">+    in_splits=args.tile_factor,
+    out_splits=4*args.tile_factor,
+    linear_cls=mpu.ColumnParallelLinear,
</span>     gather_output=False,
     init_method=init_method,
     skip_bias_add=True)
</code></span></span>

请注意我们使用in_splits和out_splits按比例缩放input_size和output_size。这会产生固定大小的图块[hidden/tile_factor, hidden/tile_factor]。

注册外部参数固定链接

已弃用 DeepSpeed 版本0.3.15引入了自动外部参数注册不再需要此步骤。

提取权重固定链接

如果您需要从 Deepspeed 中取出预训练权重您可以通过以下方式获取 fp16 权重

在 ZeRO-2 下state_dict包含 fp16 模型权重这些可以正常保存torch.save。
在 ZeRO-3 下state_dict仅包含占位符因为模型权重是跨多个 GPU 分区的。如果你想获得这些权重请启用

<span style="background-color:#263238"><span style="color:#eeffff"><code>    <span style="color:#eeffff">"zero_optimization"</span><span style="color:#eeffff">:</span> <span style="color:#eeffff">{</span>
        <span style="color:#eeffff">"stage3_gather_16bit_weights_on_model_save"</span><span style="color:#eeffff">:</span> <span style="color:#c792ea">true</span>
    <span style="color:#eeffff">}</span><span style="color:#f07178">,</span>
</code></span></span>

然后使用以下方法保存模型

<span style="background-color:#263238"><span style="color:#eeffff"><code>            <span style="color:#c792ea">if</span> <span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">deepspeed</span><span style="color:#eeffff">:</span>
                <span style="color:#eeffff">self</span><span style="color:#eeffff">.</span><span style="color:#eeffff">deepspeed</span><span style="color:#eeffff">.</span><span style="color:#eeffff">save_16bit_model</span><span style="color:#eeffff">(</span><span style="color:#eeffff">output_dir</span><span style="color:#eeffff">,</span> <span style="color:#eeffff">output_file</span><span style="color:#eeffff">)</span>
</code></span></span>

因为它需要在一个 GPU 上整合权重所以它可能很慢并且需要内存所以只在需要时使用此功能。

请注意如果stage3_gather_16bit_weights_on_model_save是False则不会保存任何权重同样因为state_dict没有它们。您也可以使用此方法保存 ZeRO-2 权重。

如果您想获得 fp32 权重我们提供了一个可以进行离线合并的特殊脚本。它不需要配置文件或 GPU。下面是它的用法示例

<span style="background-color:#263238"><span style="color:#eeffff"><code><span style="color:#f07178">$ </span><span style="color:#eeffff">cd</span> /path/to/checkpoint_dir
<span style="color:#f07178">$ </span>./zero_to_fp32.py <span style="color:#eeffff">.</span> pytorch_model.bin
Processing zero checkpoint at global_step1
Detected checkpoint of <span style="color:#eeffff">type </span>zero stage 3, world_size: 2
Saving fp32 state dict to pytorch_model.bin <span style="color:#89ddff">(</span><span style="color:#f07178">total_numel</span><span style="color:#89ddff">=</span>60506624<span style="color:#89ddff">)</span>
</code></span></span>

zero_to_fp32.py保存检查点时会自动创建脚本。

注意目前此脚本使用最终检查点大小的 2 倍内存通用 RAM。

或者如果您有足够的空闲 CPU 内存而不是获取您希望将模型更新到其 fp32 权重的文件您可以在训练结束时执行以下操作

<span style="background-color:#263238"><span style="color:#eeffff"><code>    <span style="color:#89ddff">from</span> <span style="color:#ffcb6b">deepspeed.utils.zero_to_fp32</span> <span style="color:#89ddff">import</span> <span style="color:#eeffff">load_state_dict_from_zero_checkpoint</span>
    <span style="color:#eeffff">fp32_model</span> <span style="color:#89ddff">=</span> <span style="color:#eeffff">load_state_dict_from_zero_checkpoint</span><span style="color:#eeffff">(</span><span style="color:#eeffff">deepspeed</span><span style="color:#eeffff">.</span><span style="color:#eeffff">module</span><span style="color:#eeffff">,</span> <span style="color:#eeffff">checkpoint_dir</span><span style="color:#eeffff">)</span>
</code></span></span>

请注意该模型适合保存但不再适合继续训练需要deepspeed.initialize()重新训练。

如果你只想要state_dict你可以这样做

<span style="background-color:#263238"><span style="color:#eeffff"><code>    <span style="color:#89ddff">from</span> <span style="color:#ffcb6b">deepspeed.utils.zero_to_fp32</span> <span style="color:#89ddff">import</span> <span style="color:#eeffff">get_fp32_state_dict_from_zero_checkpoint</span>
    <span style="color:#eeffff">state_dict</span> <span style="color:#89ddff">=</span> <span style="color:#eeffff">get_fp32_state_dict_from_zero_checkpoint</span><span style="color:#eeffff">(</span><span style="color:#eeffff">checkpoint_dir</span><span style="color:#eeffff">)</span>
</code></span></span>

恭喜您已完成零教程。

阿里云国内75折回扣微信号：monov8

阿里云国际，腾讯云国际，低至75折。AWS 93折免费开户实名账号代冲值优惠多多微信号：monov8 飞机：@monov6

返回列表

上一篇：The baby-bust economy “婴儿荒”经济 | 经济学人20230603版社论双语精翻

下一篇：Python中模块的动态导入和自动安装