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 通过 2 种方式集成 DeepSpeed

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

已集成什么？

训练

Accelerate 集成了 DeepSpeed ZeRO 的所有功能。这包括所有 ZeRO 阶段 1、2 和 3，以及 ZeRO-Offload、ZeRO-Infinity（可卸载到磁盘/NVMe）和 ZeRO++。下面是使用 ZeRO 进行数据并行训练的简要描述——零冗余优化器，以及来自这篇博客文章的图示。

(来源：链接)

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

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

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

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

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

f. 分层分片：通过节点内 ZeRO-3 分片和节点间数据并行训练，实现高效的多节点训练，构建于 ZeRO 阶段 3 之上。

注意：关于磁盘卸载，为了获得不错的速度，磁盘应该是 NVME，但它在技术上适用于任何磁盘。

推理

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

它是如何工作的？

先决条件：安装 DeepSpeed 版本 >=0.6.5。有关更多信息，请参阅 DeepSpeed 安装详情。

我们将首先介绍通过 accelerate config 进行的易于使用的集成。然后是更灵活、功能更丰富的 deepspeed config file 集成。

Accelerate DeepSpeed 插件

在您的机器上运行

accelerate config

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

accelerate launch my_script.py --args_to_my_script

例如，以下是使用 DeepSpeed 插件运行 NLP 示例 examples/nlp_example.py（从存储库根目录开始）的方法：

ZeRO 阶段-2 DeepSpeed 插件示例

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

accelerate launch examples/nlp_example.py --mixed_precision fp16

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

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

accelerate launch examples/nlp_example.py --mixed_precision fp16

当前，Accelerate 支持通过 CLI 进行以下配置：

`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, e.g. `pdsh`, `standard`, `openmpi`, `mvapich`, `mpich`, `slurm`, `nossh` (requires DeepSpeed >= 0.14.5). 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 配置文件

在您的机器上运行

accelerate config

并回答提出的问题。它会询问您是否要使用 DeepSpeed 的配置文件，对此您回答“是”，并提供 DeepSpeed 配置文件的路径。这将生成一个配置文件，该文件将在执行以下操作时自动用于正确设置默认选项：

accelerate launch my_script.py --args_to_my_script

例如，以下是如何使用 DeepSpeed 配置文件运行 NLP 示例 examples/by_feature/deepspeed_with_config_support.py（从存储库根目录开始）：

ZeRO 阶段-2 DeepSpeed 配置文件示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/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 的内容为：

{
    "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
}

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 配置文件示例

compute_environment: LOCAL_MACHINE
deepspeed_config:
 deepspeed_config_file: /home/ubuntu/accelerate/examples/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 的内容为：

{
    "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
}

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++ 教程修改配置文件的方法：

{
    "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 中显示此内容的片段：
```
 # 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 方法根据提供给 prepare 方法的模型、数据加载器、dummy 优化器和 dummy 调度器自动处理。只有上面示例中指定的 auto 字段由 prepare 方法处理，其余字段必须由用户明确指定。记住，前两个配置条目是通信缓冲区——它们越大，通信效率越高，占用的 GPU 内存也越多，因此这是一个可调的性能权衡。

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

为了让 auto 功能在这些 3 个配置条目上生效，Accelerate 将使用 model.config.hidden_size 或 max(model.config.hidden_sizes) 作为 hidden_size。如果两者都不可用，启动将失败，您必须手动设置这 3 个配置条目。请记住，前两个配置条目是通信缓冲区——它们越大，通信效率越高，占用的 GPU 内存也越多，因此这是一个可调的性能权衡。

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

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

代码 test.py

from accelerate import Accelerator
from accelerate.state import AcceleratorState


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


if __name__ == "__main__":
    main()

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

accelerate 配置文件的内容

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:

{
    "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 的输出

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 配置，并检查是否不再出现歧义错误。

运行 accelerate config

$ 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 配置文件的内容

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 的输出

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：将 accelerate launch 命令中与 DeepSpeed 相关的参数在 DeepSpeed 配置文件中设置为 "auto"，并检查各项是否按预期工作。

新的 ds_config.json，其中 accelerate launch DeepSpeed 命令参数为 "auto"

{
    "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 的输出

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 有 2 个选项：

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 中显示此内容的片段：

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 中显示此内容的片段：

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。以下是其用法示例：

$ 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 位模型，您可以执行：

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 感兴趣，可以执行以下操作：

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)

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

ZeRO 推理

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

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

一些需要注意的事项

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

多节点 DeepSpeed

DeepSpeed 支持通过各种启动器进行多节点推理和训练。您可以通过在 CLI 或 DeepSpeed 配置文件中设置 deepspeed_multinode_launcher 配置来指定不同的启动器。

目前，accelerate 支持为以下 DeepSpeed 多节点启动器传递配置：pdsh (默认)、standard、openmpi、mvapich、mpich、slurm、nossh（需要 DeepSpeed >= 0.14.5）。

有关不同启动器的更多信息，请阅读 DeepSpeed 文档。默认情况下，DeepSpeed 将尝试从主机器节点到其他节点使用无密码 SSH 来执行启动器命令。在此配置中，accelerate launch 命令只需在主节点上运行。如果使用 nossh 启动器，您需要在每个节点上运行 accelerate launch 命令，并使用复制的配置。

DeepSpeed 资源

与 deepspeed 相关的内部文档可以在此处找到。

论文

最后，请记住 Accelerate 只集成 DeepSpeed，因此如果您对 DeepSpeed 的使用有任何问题或疑问，请在 DeepSpeed GitHub 上提交一个 issue。

对于那些对 FSDP 和 DeepSpeed 之间的相似之处和区别感兴趣的人，请在此处的概念指南中查看！

在 GitHub 上更新