Hugging Face's logo

Accelerate 文档

DeepSpeed

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始入门

DeepSpeed

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

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

ZeRO-Offload 有其专门的论文：ZeRO-Offload：普及十亿级规模模型训练。NVMe 支持在论文 ZeRO-Infinity：打破极限规模深度学习的 GPU 内存墙 中进行了描述。

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

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

Accelerate 通过 2 个选项集成了 DeepSpeed

  1. 通过 accelerate config 中的 deepspeed config file 规范集成 DeepSpeed 功能。您只需提供自定义配置文件或使用我们的模板。本文档的大部分内容都侧重于此功能。这支持 DeepSpeed 的所有核心功能，并为用户提供了很大的灵活性。用户可能需要根据配置更改几行代码。
  2. 通过 deepspeed_plugin 集成。这支持 DeepSpeed 功能的子集，并对其余配置使用默认选项。用户无需更改任何代码，对于那些对 DeepSpeed 的大多数默认设置感到满意的用户来说，这是一个不错的选择。

集成了什么？

训练

  1. Accelerate 集成了 DeepSpeed ZeRO 的所有功能。这包括所有 ZeRO 阶段 1、2 和 3，以及 ZeRO-Offload、ZeRO-Infinity（可以将数据卸载到磁盘/NVMe）和 ZeRO++。以下是使用 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 才能获得不错的速度，但从技术上讲，它可以在任何磁盘上工作

推理

  1. DeepSpeed ZeRO 推理支持具有 ZeRO-Infinity 的 ZeRO 阶段 3。它使用与训练相同的 ZeRO 协议，但它不使用优化器和 lr 调度器，只有阶段 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（从 repo 的根目录）

ZeRO Stage-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

带有 CPU 卸载的 ZeRO Stage-3 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（从 repo 的根目录）

ZeRO Stage-2 DeepSpeed 配置文件示例

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 的内容为

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

带有 CPU 卸载的 ZeRO Stage-3 DeepSpeed 配置文件示例

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 的内容为

{
    "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 Stage 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 配置文件时的重要代码更改

  1. DeepSpeed 优化器和调度器。有关这些的更多信息，请参阅 DeepSpeed 优化器 和 DeepSpeed 调度器 文档。我们将了解在使用这些时代码中需要的更改。

    a. DS Optim + DS Scheduler：optimizer 和 scheduler 键都存在于 DeepSpeed 配置文件中的情况。在这种情况下，将使用这些键，用户必须使用 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. 自定义 Optim + 自定义 Scheduler：optimizer 和 scheduler 键在 DeepSpeed 配置文件中都不存在的情况。在这种情况下，用户无需进行代码更改，这是通过 DeepSpeed 插件集成时的情况。在上面的示例中，我们可以看到，如果 optimizer 和 scheduler 键在 DeepSpeed 配置文件中不存在，则代码保持不变。

    c. 自定义 Optim + DS Scheduler：只有 scheduler 键存在于 DeepSpeed 配置文件中的情况。在这种情况下，用户必须使用 accelerate.utils.DummyScheduler 来替换其代码中的 PyTorch/自定义调度器。

    d. DS Optim + 自定义 Scheduler：只有 optimizer 键存在于 DeepSpeed 配置文件中的情况。这将导致错误，因为只有在使用 DS Optim 时才能使用 DS Scheduler。

  2. 请注意上面示例 DeepSpeed 配置文件中的 auto 值。这些值由 prepare 方法根据提供给 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

为了使 auto 功能适用于这 3 个配置条目 - Accelerate 将使用 model.config.hidden_size 或 max(model.config.hidden_sizes) 作为 hidden_size。如果两者都不可用，则启动将失败，您将不得不手动设置这 3 个配置条目。请记住，前 2 个配置条目是通信缓冲区 - 它们越大，通信效率越高，它们越大，它们将消耗的 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：手动篡改了包含 deepspeed_config_file 以及其他条目的 accelerate 配置文件。

  1. 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
  1. 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
}
  1. 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 配置，并检查现在是否未抛出歧义错误。

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

  1. 新的 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
}
  1. 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}}

注意:

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

保存和加载

  1. ZeRO Stage-1 和 Stage-2 的模型保存和加载保持不变。

  2. 在 ZeRO Stage-3 下，state_dict 仅包含占位符，因为模型权重已跨多个 GPU 分区。ZeRO Stage-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 stage 3。它使用与训练相同的 ZeRO 协议，但它不使用优化器和 lr 调度器，并且只有 stage 3 是相关的。通过 accelerate 集成，您只需如下所示准备模型和数据加载器即可

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

需要注意的一些注意事项

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

多节点 DeepSpeed

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

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

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

DeepSpeed 资源

有关 deepspeed 内部原理的文档可以在这里找到。

论文

最后，请记住 Accelerate 仅集成了 DeepSpeed，因此如果您在使用 DeepSpeed 方面有任何问题或疑问，请在 DeepSpeed GitHub 上提交 issue。

对于那些对 FSDP 和 DeepSpeed 之间的异同感兴趣的人，请查看此处的概念指南！

< > 在 GitHub 上更新

←低精度 (FP8) 训练 使用 DeepSpeed 的多个模型→