DeepSpeed

DeepSpeed 实现了 ZeRO 论文中描述的所有内容。一些重要的优化包括：

优化器状态分区（ZeRO 阶段 1）
梯度分区（ZeRO 阶段 2）
参数分区（ZeRO 阶段 3）
自定义混合精度训练处理
一系列基于快速 CUDA 扩展的优化器
ZeRO-Offload 到 CPU 和磁盘/NVMe
模型参数的分层分区（ZeRO++）

ZeRO-Offload 有其专门的论文：ZeRO-Offload: 民主化数十亿规模模型的训练。而 NVMe 支持在论文 ZeRO-Infinity: 打破 GPU 内存墙以实现大规模深度学习中有描述。

DeepSpeed ZeRO-2 主要仅用于训练，因为其功能对推理无用。

DeepSpeed ZeRO-3 也可以用于推理，因为它允许在多个 GPU 上加载巨大的模型，这在单个 GPU 上是不可能的。

Accelerate 通过以下两种选项集成 DeepSpeed：

通过 accelerate config 中的 deepspeed 配置文件 规范集成 DeepSpeed 功能。你只需提供自定义配置文件或使用我们的模板。本文档的大部分内容都集中在这一功能上。这支持 DeepSpeed 的所有核心功能，并为用户提供了很大的灵活性。根据配置，用户可能需要更改几行代码。
通过 deepspeed_plugin 集成。这支持 DeepSpeed 的部分功能，并为其余配置使用默认选项。用户无需更改任何代码，适用于那些对 DeepSpeed 的大多数默认设置满意的人。

集成了什么？

训练：

Accelerate 集成了 DeepSpeed ZeRO 的所有功能。这包括所有 ZeRO 阶段 1、2 和 3 以及 ZeRO-Offload、ZeRO-Infinity（可以卸载到磁盘/NVMe）和 ZeRO++。以下是使用 ZeRO - Zero Redundancy Optimizer 的数据并行性的简要描述，以及来自这篇博客文章的图表：

(来源: 链接)

a. 阶段 1 : 在数据并行工作器/GPU 之间分片优化器状态

b. 阶段 2 : 在数据并行工作器/GPU 之间分片优化器状态 + 梯度

c. 阶段 3: 在数据并行工作器/GPU 之间分片优化器状态 + 梯度 + 模型参数

d. 优化器卸载: 在 ZeRO 阶段 2 的基础上，将梯度 + 优化器状态卸载到 CPU/磁盘

e. 参数卸载: 在 ZeRO 阶段 3 的基础上，将模型参数卸载到 CPU/磁盘

f. 分层分区: 在 ZeRO 阶段 3 的基础上，支持节点间的数据并行训练和节点内的 ZeRO-3 分片，实现高效的多节点训练

Note

关于磁盘卸载，磁盘应该是 NVMe 以获得良好的速度，但技术上它可以在任何磁盘上工作。

推理：

DeepSpeed ZeRO 推理支持 ZeRO 阶段 3 与 ZeRO-Infinity。它使用与训练相同的 ZeRO 协议，但不使用优化器和学习率调度器，只有阶段 3 是相关的。更多详情请参见： deepspeed-zero-inference。

它是如何工作的？

先决条件：安装 DeepSpeed 版本 >=0.6.5。请参阅 DeepSpeed 安装详情以获取更多信息。

我们首先看看通过 accelerate config 进行的简单集成。然后是更灵活且功能丰富的 deepspeed config file 集成。

Accelerate DeepSpeed 插件

在你的机器上只需运行：

bash

accelerate config

并回答提出的问题。它会问你是否要使用 DeepSpeed 的配置文件，你应该回答“不”。然后回答以下问题以生成一个基本的 DeepSpeed 配置。这将生成一个配置文件，该文件将自动用于在执行时正确设置默认选项。

bash

accelerate launch my_script.py --args_to_my_script

例如，你可以这样运行 NLP 示例 examples/nlp_example.py（从仓库的根目录）使用 DeepSpeed 插件：

ZeRO 阶段 2 DeepSpeed 插件示例

bash

compute_environment: LOCAL_MACHINE
deepspeed_config:
 gradient_accumulation_steps: 1
 gradient_clipping: 1.0
 offload_optimizer_device: none
 offload_param_device: none
 zero3_init_flag: true
 zero_stage: 2
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

bash

accelerate launch examples/nlp_example.py --mixed_precision fp16

ZeRO 阶段 3 与 CPU 卸载 DeepSpeed 插件示例

在这个示例中，我们将展示如何使用 DeepSpeed 的 ZeRO 阶段 3 和 CPU 卸载功能来训练大型模型。ZeRO 阶段 3 通过将优化器状态和梯度卸载到 CPU 内存，从而显著减少了 GPU 内存的使用。这使得在单个 GPU 上训练更大的模型成为可能。

准备工作

安装 DeepSpeed：确保你已经安装了 DeepSpeed。你可以使用以下命令进行安装：
bash
```
pip install deepspeed
```
准备数据集：确保你有一个可用的数据集。你可以使用 Hugging Face 的 datasets 库来加载数据集：
python
```
from datasets import load_dataset

dataset = load_dataset('your_dataset_name')
```

配置 DeepSpeed 配置文件：创建一个 DeepSpeed 配置文件 ds_config.json，内容如下：

json

{
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 3,
    "offload_optimizer": {
      "device": "cpu",
      "pin_memory": true
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "overlap_comm": true,
    "load_from_fp32_weights": true,
    "elastic_checkpoint": true
  }
}

训练模型

导入必要的库：

python

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from deepspeed import DeepSpeedEngine, DeepSpeedConfig

加载模型和分词器：

python

model_name = 'your_model_name'
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

初始化 DeepSpeed：

python

ds_config = DeepSpeedConfig('ds_config.json')
model, optimizer, _, _ = DeepSpeedEngine(model, ds_config)

定义训练循环：

python

for epoch in range(num_epochs):
    for batch in dataloader:
        inputs = tokenizer(batch['text'], return_tensors='pt', padding=True, truncation=True).to('cuda')
        outputs = model(**inputs, labels=inputs['input_ids'])
        loss = outputs.loss
        model.backward(loss)
        model.step()

总结

通过使用 DeepSpeed 的 ZeRO 阶段 3 和 CPU 卸载功能，我们可以在单个 GPU 上训练更大的模型，而不会遇到内存不足的问题。这为训练大型模型提供了极大的灵活性和效率。

bash

compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: cpu
  offload_param_device: cpu
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

bash

accelerate launch examples/nlp_example.py --mixed_precision fp16

目前，Accelerate 通过 CLI 支持以下配置：

bash

`zero_stage`: [0] Disabled, [1] optimizer state partitioning, [2] optimizer+gradient state partitioning and [3] optimizer+gradient+parameter partitioning
`gradient_accumulation_steps`: Number of training steps to accumulate gradients before averaging and applying them.
`gradient_clipping`: Enable gradient clipping with value.
`offload_optimizer_device`: [none] Disable optimizer offloading, [cpu] offload optimizer to CPU, [nvme] offload optimizer to NVMe SSD. Only applicable with ZeRO >= Stage-2.
`offload_optimizer_nvme_path`: Decides Nvme Path to offload optimizer states. If unspecified, will default to 'none'.
`offload_param_device`: [none] Disable parameter offloading, [cpu] offload parameters to CPU, [nvme] offload parameters to NVMe SSD. Only applicable with ZeRO Stage-3.
`offload_param_nvme_path`: Decides Nvme Path to offload parameters. If unspecified, will default to 'none'.
`zero3_init_flag`: Decides whether to enable `deepspeed.zero.Init` for constructing massive models. Only applicable with ZeRO Stage-3.
`zero3_save_16bit_model`: Decides whether to save 16-bit model weights when using ZeRO Stage-3.
`mixed_precision`: `no` for FP32 training, `fp16` for FP16 mixed-precision training and `bf16` for BF16 mixed-precision training.
`deepspeed_moe_layer_cls_names`: Comma-separated list of transformer Mixture-of-Experts (MoE) layer class names (case-sensitive) to wrap ,e.g, `MixtralSparseMoeBlock`, `Qwen2MoeSparseMoeBlock`, `JetMoEAttention,JetMoEBlock` ...
`deepspeed_hostfile`: DeepSpeed hostfile for configuring multi-node compute resources.
`deepspeed_exclusion_filter`: DeepSpeed exclusion filter string when using mutli-node setup.
`deepspeed_inclusion_filter`: DeepSpeed inclusion filter string when using mutli-node setup.
`deepspeed_multinode_launcher`: DeepSpeed multi-node launcher to use. If unspecified, will default to `pdsh`.
`deepspeed_config_file`: path to the DeepSpeed config file in `json` format. See the next section for more details on this.

要调整更多选项，你需要使用一个 DeepSpeed 配置文件。

DeepSpeed 配置文件

在你的机器上运行：

bash

accelerate config

并回答提出的问题。它会问你是否要使用 deepspeed 的配置文件，你回答“是” 并提供 deepspeed 配置文件的路径。这将生成一个配置文件，该文件将自动用于在执行时正确设置默认选项。

bash

accelerate launch my_script.py --args_to_my_script

例如，你可以这样运行 NLP 示例 examples/by_feature/deepspeed_with_config_support.py（从仓库的根目录）并使用 DeepSpeed 配置文件：

ZeRO 阶段 2 DeepSpeed 配置文件示例

bash

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage2_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

zero_stage2_config.json 的内容为：

json

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto",
            "torch_adam": true,
            "adam_w_mode": true
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

bash

accelerate launch examples/by_feature/deepspeed_with_config_support.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "./clm/clm_deepspeed_stage2_accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 24 \
--per_device_eval_batch_size 24 \
--num_train_epochs 3 \
--with_tracking \
--report_to "wandb"\

ZeRO 阶段 3 带 CPU 卸载的 DeepSpeed 配置文件示例

bash

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/configs/deepspeed_config_templates/zero_stage3_offload_config.json
 zero3_init_flag: true
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
mixed_precision: fp16
num_machines: 1
num_processes: 2
use_cpu: false

zero_stage3_offload_config.json 的内容为：

json

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "weight_decay": "auto"
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto",
            "total_num_steps": "auto"
        }
    },
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "sub_group_size": 1e9,
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": "auto"
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

bash

accelerate launch examples/by_feature/deepspeed_with_config_support.py \
--config_name "gpt2-large" \
--tokenizer_name "gpt2-large" \
--dataset_name "wikitext" \
--dataset_config_name "wikitext-2-raw-v1" \
--block_size 128 \
--output_dir "./clm/clm_deepspeed_stage3_offload_accelerate" \
--learning_rate 5e-4 \
--per_device_train_batch_size 32 \
--per_device_eval_batch_size 32 \
--num_train_epochs 3 \
--with_tracking \
--report_to "wandb"\

ZeRO++ 配置示例

你可以通过使用适当的配置参数来使用 ZeRO++ 的功能。请注意，ZeRO++ 是 ZeRO 阶段 3 的扩展。以下是如何修改配置文件，来自 DeepSpeed 的 ZeRO++ 教程：

json

{
    "zero_optimization": {
        "stage": 3,
        "reduce_bucket_size": "auto",

        "zero_quantized_weights": true,
        "zero_hpz_partition_size": 8,
        "zero_quantized_gradients": true,

        "contiguous_gradients": true,
        "overlap_comm": true
    }
}

对于分层分区，分区大小 zero_hpz_partition_size 最好设置为每个节点的 GPU 数量。（例如，上述配置文件假设每个节点有 8 个 GPU）

使用 DeepSpeed 配置文件时的重要代码更改

DeepSpeed 优化器和调度器。有关这些的更多信息，请参阅 DeepSpeed 优化器和 DeepSpeed 调度器文档。我们将看看在使用这些时代码中需要的更改。
a. DS 优化器 + DS 调度器：当 DeepSpeed 配置文件中同时存在 optimizer 和 scheduler 键时的情况。在这种情况下，将使用这些配置，用户需要使用 accelerate.utils.DummyOptim 和 accelerate.utils.DummyScheduler 来替换代码中的 PyTorch/自定义优化器和调度器。以下是 examples/by_feature/deepspeed_with_config_support.py 中的代码片段，展示了这一点：

python

    # Creates Dummy Optimizer if `optimizer` was specified in the config file else creates Adam Optimizer
    optimizer_cls = (
        torch.optim.AdamW
        if accelerator.state.deepspeed_plugin is None
        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
        else DummyOptim
    )
    optimizer = optimizer_cls(optimizer_grouped_parameters, lr=args.learning_rate)

    # Creates Dummy Scheduler if `scheduler` was specified in the config file else creates `args.lr_scheduler_type` Scheduler
    if (
        accelerator.state.deepspeed_plugin is None
        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
    ):
        lr_scheduler = get_scheduler(
            name=args.lr_scheduler_type,
            optimizer=optimizer,
            num_warmup_steps=args.num_warmup_steps,
            num_training_steps=args.max_train_steps,
        )
    else:
        lr_scheduler = DummyScheduler(
            optimizer, total_num_steps=args.max_train_steps, warmup_num_steps=args.num_warmup_steps
        )

b. 自定义优化器 + 自定义调度器：当 DeepSpeed 配置文件中缺少 optimizer 和 scheduler 键时的情况。在这种情况下，用户不需要进行任何代码更改，这是通过 DeepSpeed 插件进行集成时的情况。在上面的示例中，我们可以看到，如果 DeepSpeed 配置文件中缺少 optimizer 和 scheduler 键，代码将保持不变。

c. 自定义优化器 + DS 调度器：当 DeepSpeed 配置文件中仅存在 scheduler 键时的情况。在这种情况下，用户需要使用 accelerate.utils.DummyScheduler 替换代码中的 PyTorch/自定义调度器。

d. DS 优化器 + 自定义调度器：当 DeepSpeed 配置文件中仅存在 optimizer 键时的情况。这将导致错误，因为只有在使用 DS 优化器时才能使用 DS 调度器。

注意上述示例 DeepSpeed 配置文件中的 auto 值。这些值由 prepare 方法根据提供的模型、数据加载器、虚拟优化器和虚拟调度器自动处理。只有上述示例中指定的 auto 字段由 prepare 方法处理，其余字段需要用户显式指定。

auto 值的计算方式如下：

reduce_bucket_size: hidden_size * hidden_size
stage3_prefetch_bucket_size: int(0.9 * hidden_size * hidden_size)
stage3_param_persistence_threshold: 10 * hidden_size

为了使这 3 个配置项的 auto 功能生效，Accelerate 将使用 model.config.hidden_size 或 max(model.config.hidden_sizes) 作为 hidden_size。如果这些值都不可用，启动将失败，你将需要手动设置这 3 个配置项。请记住，前两个配置项是通信缓冲区，它们越大，通信效率越高，但同时也会消耗更多的 GPU 内存，因此这是一个可调的性能权衡。

使用 DeepSpeed 配置文件时需要注意的事项

以下是使用 deepspeed_config_file 在不同场景下的示例脚本。

代码 test.py：

python

from accelerate import Accelerator
from accelerate.state import AcceleratorState


def main():
    accelerator = Accelerator()
    accelerator.print(f"{AcceleratorState()}")


if __name__ == "__main__":
    main()

场景 1：手动修改的加速配置文件，包含 deepspeed_config_file 以及其他条目。

accelerate 配置文件的内容：

yaml

command_file: null
commands: null
compute_environment: LOCAL_MACHINE
deepspeed_config:
  gradient_accumulation_steps: 1
  gradient_clipping: 1.0
  offload_optimizer_device: 'cpu'
  offload_param_device: 'cpu'
  zero3_init_flag: true
  zero3_save_16bit_model: true
  zero_stage: 3
  deepspeed_config_file: 'ds_config.json'
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
gpu_ids: null
machine_rank: 0
main_process_ip: null
main_process_port: null
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 2
rdzv_backend: static
same_network: true
tpu_name: null
tpu_zone: null
use_cpu: false

ds_config.json:

json

{
    "bf16": {
        "enabled": true
    },
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": false,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_clipping": 1.0,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": 10,
    "steps_per_print": 2000000
}

accelerate launch test.py 的输出：

bash

ValueError: When using `deepspeed_config_file`, the following accelerate config variables will be ignored:
['gradient_accumulation_steps', 'gradient_clipping', 'zero_stage', 'offload_optimizer_device', 'offload_param_device',
'zero3_save_16bit_model', 'mixed_precision'].
Please specify them appropriately in the DeepSpeed config file.
If you are using an accelerate config file, remove other config variables mentioned in the above specified list.
The easiest method is to create a new config following the questionnaire via `accelerate config`.
It will only ask for the necessary config variables when using `deepspeed_config_file`.

场景 2：使用错误的解决方案创建新的加速配置，并检查现在是否不再抛出歧义错误。

运行 accelerate config：

bash

$ accelerate config
-------------------------------------------------------------------------------------------------------------------------------
In which compute environment are you running?
This machine
-------------------------------------------------------------------------------------------------------------------------------
Which type of machine are you using?
multi-GPU
How many different machines will you use (use more than 1 for multi-node training)? [1]:
Do you wish to optimize your script with torch dynamo?[yes/NO]:
Do you want to use DeepSpeed? [yes/NO]: yes
Do you want to specify a json file to a DeepSpeed config? [yes/NO]: yes
Please enter the path to the json DeepSpeed config file: ds_config.json
Do you want to enable `deepspeed.zero.Init` when using ZeRO Stage-3 for constructing massive models? [yes/NO]: yes
How many GPU(s) should be used for distributed training? [1]:4
accelerate configuration saved at ds_config_sample.yaml

accelerate 配置的内容：

yaml

compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_config_file: ds_config.json
  zero3_init_flag: true
distributed_type: DEEPSPEED
downcast_bf16: 'no'
dynamo_backend: 'NO'
fsdp_config: {}
machine_rank: 0
main_training_function: main
megatron_lm_config: {}
num_machines: 1
num_processes: 4
rdzv_backend: static
same_network: true
use_cpu: false

accelerate launch test.py 的输出：

bash

Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: bf16
ds_config: {'bf16': {'enabled': True}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': False, 'offload_optimizer': {'device': 'none'}, 'offload_param': {'device': 'none'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 10, 'steps_per_print': inf, 'fp16': {'enabled': False}}

场景 3：在 DeepSpeed 配置文件中将 accelerate launch 命令的 DeepSpeed 参数设置为 "auto"，并检查一切是否按预期工作。

新的 ds_config.json 文件，将 accelerate launch DeepSpeed 命令参数设置为 "auto"：

json

{
    "bf16": {
        "enabled": "auto"
    },
    "zero_optimization": {
        "stage": "auto",
        "stage3_gather_16bit_weights_on_model_save": "auto",
        "offload_optimizer": {
            "device": "auto"
        },
        "offload_param": {
            "device": "auto"
        }
    },
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "steps_per_print": 2000000
}

accelerate launch --mixed_precision="fp16" --zero_stage=3 --gradient_accumulation_steps=5 --gradient_clipping=1.0 --offload_param_device="cpu" --offload_optimizer_device="nvme" --zero3_save_16bit_model="true" test.py 的输出：

bash

Distributed environment: DEEPSPEED  Backend: nccl
Num processes: 4
Process index: 0
Local process index: 0
Device: cuda:0
Mixed precision type: fp16
ds_config: {'bf16': {'enabled': False}, 'zero_optimization': {'stage': 3, 'stage3_gather_16bit_weights_on_model_save': True, 'offload_optimizer': {'device': 'nvme'}, 'offload_param': {'device': 'cpu'}}, 'gradient_clipping': 1.0, 'train_batch_size': 'auto', 'train_micro_batch_size_per_gpu': 'auto', 'gradient_accumulation_steps': 5, 'steps_per_print': inf, 'fp16': {'enabled': True, 'auto_cast': True}}

注意:

剩余的 "auto" 值在 accelerator.prepare() 调用中处理，如 使用 DeepSpeed 配置文件时的重要代码更改 中的第 2 点所述。
仅当 gradient_accumulation_steps 为 auto 时，通过 Accelerator(gradient_accumulation_steps=k) 创建 Accelerator 对象时传递的值才会被使用。使用 DeepSpeed 插件时，将使用插件中的值，并覆盖创建 Accelerator 对象时传递的值。

保存和加载

ZeRO 阶段 1 和阶段 2 的模型保存和加载方式保持不变。
在 ZeRO 阶段 3 下，state_dict 仅包含占位符，因为模型权重分布在多个 GPU 上。 ZeRO 阶段 3 有两个选项：
a. 将整个 16 位模型权重保存，以便稍后使用 model.load_state_dict(torch.load(pytorch_model.bin)) 直接加载。为此，可以在 DeepSpeed 配置文件中将 zero_optimization.stage3_gather_16bit_weights_on_model_save 设置为 True，或在 DeepSpeed 插件中将 zero3_save_16bit_model 设置为 True。 请注意，此选项需要在单个 GPU 上整合权重，可能会很慢且占用大量内存，因此仅在需要时使用此功能。 以下是 examples/by_feature/deepspeed_with_config_support.py 中的示例代码片段：

python

   unwrapped_model = accelerator.unwrap_model(model)

   # New Code #
   # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
   # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
   # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
   # For Zero Stages 1 and 2, models are saved as usual in the output directory.
   # The model name saved is `pytorch_model.bin`
   unwrapped_model.save_pretrained(
       args.output_dir,
       is_main_process=accelerator.is_main_process,
       save_function=accelerator.save,
       state_dict=accelerator.get_state_dict(model),
   )

b. 要获取 32 位权重，首先使用 model.save_checkpoint() 保存模型。以下是来自 examples/by_feature/deepspeed_with_config_support.py 的代码片段，展示了这一点：

python

   success = model.save_checkpoint(PATH, ckpt_id, checkpoint_state_dict)
   status_msg = f"checkpointing: PATH={PATH}, ckpt_id={ckpt_id}"
   if success:
       logging.info(f"Success {status_msg}")
   else:
       logging.warning(f"Failure {status_msg}")

这将在检查点目录中创建 ZeRO 模型和优化器分区，以及 zero_to_fp32.py 脚本。你可以使用此脚本来进行离线整合。它不需要任何配置文件或 GPU。以下是一个使用示例：

bash

   $ cd /path/to/checkpoint_dir
   $ ./zero_to_fp32.py . pytorch_model.bin
   Processing zero checkpoint at global_step1
   Detected checkpoint of type zero stage 3, world_size: 2
   Saving fp32 state dict to pytorch_model.bin (total_numel=60506624)

要获取 32 位模型以进行保存/推理，你可以执行：

python

   from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint

   unwrapped_model = accelerator.unwrap_model(model)
   fp32_model = load_state_dict_from_zero_checkpoint(unwrapped_model, checkpoint_dir)

如果你只对 state_dict 感兴趣，可以这样做：

python

   from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint

   state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir)

请注意，所有这些函数都需要大约两倍于最终检查点大小的内存（通用 RAM）。

ZeRO 推理

DeepSpeed ZeRO 推理支持 ZeRO 阶段 3 和 ZeRO-Infinity。它使用与训练相同的 ZeRO 协议，但不使用优化器和学习率调度器，只有阶段 3 是相关的。通过与 accelerate 集成，你只需按照以下方式准备模型和数据加载器：

python

model, eval_dataloader = accelerator.prepare(model, eval_dataloader)

需要注意的几个问题

当前集成不支持 DeepSpeed 的管道并行。
当前集成不支持 mpu，限制了 Megatron-LM 中支持的张量并行。
当前集成不支持多个模型。

DeepSpeed 资源

有关 DeepSpeed 内部的文档可以在这里找到这里。

论文：

最后，请记住，Accelerate 仅集成 DeepSpeed，因此如果你在使用 DeepSpeed 时遇到任何问题或有疑问，请在 DeepSpeed GitHub 上提交问题。

DeepSpeed ​

集成了什么？ ​

它是如何工作的？ ​

Accelerate DeepSpeed 插件 ​

准备工作 ​

训练模型 ​

总结 ​

DeepSpeed 配置文件 ​

保存和加载 ​

ZeRO 推理 ​

需要注意的几个问题 ​

DeepSpeed 资源 ​

DeepSpeed

集成了什么？

它是如何工作的？

Accelerate DeepSpeed 插件

准备工作

训练模型

总结

DeepSpeed 配置文件

保存和加载

ZeRO 推理

需要注意的几个问题

DeepSpeed 资源