Hugging Face 的 LOGO

Transformers 文档

DeepSpeed

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

DeepSpeed

DeepSpeed 旨在通过数据、模型、流水线甚至所有这三者的组合 并行 策略来优化大型模型的分布式训练，从而提供更好的内存效率和更快的训练速度。这是通过 ZeRO（零冗余优化器） 实现的，它包含三个阶段。

ZeRO 阶段 description
1 分区优化器状态
2 分区优化器和梯度状态
3 分区优化器、梯度和参数

每个阶段逐步节省更多内存，允许非常大的模型拟合并在单个 GPU 上进行训练。所有 ZeRO 阶段，将优化器内存和计算从 GPU 卸载到 CPU，都已与 Trainer 集成。为 Trainer 提供配置文件或其中一个示例模板，即可启用 DeepSpeed 功能。

本指南将引导您完成设置 DeepSpeed 配置文件、如何在 Trainer 中启用其功能以及部署进行训练。

从 PyPI 或 Transformers 安装 DeepSpeed。有关更详细的安装说明，请参阅 DeepSpeed 安装 或 GitHUB README。

PyPI
Transformers
pip install deepspeed

如果您在安装过程中遇到问题，请参阅 DeepSpeed CUDA 安装。虽然 DeepSpeed 有一个可 pip 安装的包，但强烈建议 从源代码安装，以确保它与您的硬件匹配，并支持 PyPI 分发版中不可用的某些功能。

DeepSpeed 提供了一个工具，用于估算参数、优化器和梯度状态所需的 CPU 和 GPU 内存。您还需要为 CUDA 内核和激活预留一些内存。

运行以下命令以检查 bigscience/T0_3B 在单个 GPU 上的内存要求。

$ python -c 'from transformers import AutoModel; \
from deepspeed.runtime.zero.stage3 import estimate_zero3_model_states_mem_needs_all_live; \
model = AutoModel.from_pretrained("bigscience/T0_3B"); \
estimate_zero3_model_states_mem_needs_all_live(model, num_gpus_per_node=1, num_nodes=1)'
[...]
Estimated memory needed for params, optim states and gradients for a:
HW: Setup with 1 node, 1 GPU per node.
SW: Model with 2783M total params, 65M largest layer params.
  per CPU  |  per GPU |   Options
   70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=1
   70.00GB |   0.25GB | offload_param=cpu , offload_optimizer=cpu , zero_init=0
   62.23GB |   5.43GB | offload_param=none, offload_optimizer=cpu , zero_init=1
   62.23GB |   5.43GB | offload_param=none, offload_optimizer=cpu , zero_init=0
    0.37GB |  46.91GB | offload_param=none, offload_optimizer=none, zero_init=1
   15.56GB |  46.91GB | offload_param=none, offload_optimizer=none, zero_init=0

如果您有足够的 GPU 内存，请禁用 CPU 和 NVMe 卸载以加速所有操作。

选择 ZeRO 阶段

请参考下表，以帮助您选择合适的 ZeRO 阶段进行训练，因为训练速度和内存使用之间存在权衡。表中按 ZeRO 阶段从快到慢，从内存使用最少到最多的顺序排列。

最快 最少内存使用
ZeRO-1 ZeRO-3 + 卸载
ZeRO-2 ZeRO-3
ZeRO-2 + 卸载 ZeRO-2 + 卸载
ZeRO-3 ZeRO-2
ZeRO-3 + 卸载 ZeRO-1

确定您要优化的性能类型（速度或内存），然后倒推以找到最适合您用例的 ZeRO 阶段。例如，如果您优化速度，请从最快的 ZeRO 阶段开始，如果内存不足，请尝试下一个阶段，它速度较慢但内存效率更高。

配置文件

确定 ZeRO 阶段后，设置配置文件以通过 Trainer 启用 DeepSpeed。配置文件包含如何配置和设置训练的所有参数。执行训练脚本时，DeepSpeed 会将 Trainer 的配置记录到控制台，以便您确切地看到正在使用什么。

在 DeepSpeed Configuration JSON 参考中查找 DeepSpeed 配置选项的完整列表。在 DeepSpeedExamples 主 DeepSpeed 存储库中，也有各种 DeepSpeed 配置的实际示例。运行以下命令可快速查找特定示例。

git clone https://github.com/microsoft/DeepSpeedExamples
cd DeepSpeedExamples
find . -name '*json'
# find examples with the Lamb optimizer
grep -i Lamb $(find . -name '*json')

如果您是从命令行界面进行训练，配置文件将作为 JSON 文件的路径传递；如果您在笔记本电脑中使用 Trainer，则它将作为嵌套字典对象传递。

文件路径
嵌套字典
TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...,
)

DeepSpeed 与 Trainer 参数对比

有三种类型的配置参数。

  1. 一些配置参数由 DeepSpeed 和 Trainer 共享，当定义冲突时，这会使识别错误变得困难。在这种情况下，请通过 Trainer 命令行参数配置这些参数。
  2. 一些配置参数会自动从模型配置中派生，无需手动配置。Trainer 使用配置值 auto 来设置最正确或最高效的选项。您可以显式定义这些参数，但必须注意确保 Trainer 和 DeepSpeed 配置参数匹配。不匹配可能会导致训练失败，且极难检测。
  3. 一些配置参数是 DeepSpeed 特有的，应根据您的训练要求手动设置。

有两种修改配置参数的方法。

一些值（例如 scheduler.params.total_num_steps）会在训练期间由 Trainer 计算。

  1. 创建或加载一个 DeepSpeed 配置作为主配置。
  2. 基于 DeepSpeed 配置值创建 TrainingArguments 对象。

ZeRO 阶段

每个 ZeRO 阶段的配置都在 zero_optimization 中定义。

有关每个参数的更详细说明，请参阅 DeepSpeed Configuration JSON 参考。这些参数必须与 DeepSpeed 一起设置，因为 Trainer 不提供等效的命令行参数。

DeepSpeed 不会验证参数名称，任何拼写错误都会回退到参数的默认设置。请观察 DeepSpeed 引擎启动日志消息，以查看正在使用哪些值。

ZeRO-1
ZeRO-2
ZeRO-3

ZeRO-1 将优化器状态分片到 GPU 上，您可以期待 небольшой 速度提升。

{
    "zero_optimization": {
        "stage": 1
    }
}

NVMe

ZeRO-Infinity 将模型状态卸载到 CPU 和/或 NVMe 以节省更多内存。智能分区和分块算法允许每个 GPU 在卸载过程中发送和接收非常小量的数据，因此现代 NVMe 可以容纳比您的训练过程可用内存总量更大的内存池。ZeRO-Infinity 需要 ZeRO-3。

根据可用的 CPU 和 NVMe 内存，您可以卸载 优化器状态 和 参数，或者只卸载其中一个，或者都不卸载。确保 nvme_path 指向 NVMe 设备，因为虽然它在普通硬盘或固态硬盘上也能工作，但会慢得多。使用现代 NVMe，您可以期望峰值传输速度对于读取操作约为 3.5GB/s，对于写入操作约为 3GB/s。

考虑对您的训练设置运行 基准测试 以确定最佳的 aio 配置。

下面的示例 ZeRO-3 和 ZeRO-Infinity 配置将大多数参数值设置为 auto，但您也可以手动配置这些值。

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_num_steps": "auto"
        }
    },

    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": {
            "device": "nvme",
            "nvme_path": "/local_nvme",
            "pin_memory": true,
            "buffer_count": 4,
            "fast_init": false
        },
        "offload_param": {
            "device": "nvme",
            "nvme_path": "/local_nvme",
            "pin_memory": true,
            "buffer_count": 5,
            "buffer_size": 1e8,
            "max_in_cpu": 1e9
        },
        "overlap_comm": true,
        "contiguous_gradients": true,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "aio": {
        "block_size": 262144,
        "queue_depth": 32,
        "thread_count": 1,
        "single_submit": false,
        "overlap_events": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}

序列并行

DeepSpeed 的 ALST/Ulysses 序列并行通过将序列跨越多个 GPU 进行分割，从而实现非常长的序列的训练。这对于训练具有非常长序列长度的大型语言模型特别有用。

Arctic Long Sequence Training (ALST) 使用输入沿序列维度分片和注意力头并行的组合。通过这种方法，您可以在单个 H100 GPU 上训练序列长度高达 500K 个 token 的模型，在单个节点上训练 3.7M 个 token，或者在仅四个节点上使用 Llama-8B 训练 15M 个 token。这里描述的实现启用了完整 ALST 系统的其中一个组件。有关 TiledMLP 和激活检查点卸载等其他优化，请参阅 DeepSpeed ALST 教程。

有关序列并行的更详细信息，请参阅 Accelerate 序列并行 指南。

要使用 Trainer 启用 ALST/Ulysses 序列并行，请在 TrainingArguments 中配置 parallelism_config。序列并行通过 Accelerate 的 ParallelismConfig 进行配置，并且需要 Accelerate 版本高于 1.12.0。

from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

# Example: 4 GPUs with sp_size=4, dp_replicate_size=1 (no data parallelism)
# Ensure total_size = dp_replicate_size * dp_shard_size * sp_size = 1 * 1 * 4 = 4 GPUs
parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,  # Number of GPUs to split sequence across
    dp_replicate_size=1,  # Explicit: no data parallelism
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="sdpa",
    ),
)

training_args = TrainingArguments(
    ...,
    deepspeed="path/to/deepspeed_config.json",
    parallelism_config=parallelism_config,
)

您也可以使用 Accelerate 配置文件来配置序列并行。

distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 4  # Total number of processes
parallelism_config:
  parallelism_config_sp_size: 4  # Sequence parallel size
  parallelism_config_dp_replicate_size: 1  # Must be: dp_replicate_size * dp_shard_size * sp_size = num_processes
  parallelism_config_sp_backend: deepspeed
  parallelism_config_sp_seq_length_is_variable: true
  parallelism_config_sp_attn_implementation: sdpa

重要的配置参数包括以下内容。

  • sp_backend 必须设置为 "deepspeed" 才能使用 ALST/Ulysses 序列并行。
  • sp_size 是序列并行的度。例如，sp_size=4 表示 4 个 GPU 将并行处理单个序列。您至少需要 2 个 GPU 才能启用序列并行。数据馈送：每个 rank 从 DataLoader 接收唯一的 Dataloader 流（类似于 DP）。批大小计算：有效 dp_world_size = world_size / sp_size。因此，对于 4 个 GPU 和 sp_size=4，每个 rank 从 DataLoader 接收不同的样本，但 dp_world_size=1 用于总批大小计算。
  • sp_seq_length_is_variable 决定如何处理序列长度。当设置为 True（推荐）时，实现会适应批次之间变化的序列长度。当设置为 False 时，所有序列都必须填充到 sp_seq_length 指定的固定长度。
  • sp_attn_implementation 指定要使用的注意力实现。支持的值为 "sdpa"、"flash_attention_2" 或 "flash_attention_3"。Flash Attention 推荐用于最佳性能，尤其是在批处理中具有多个样本时，因为 SDPA 可能会错误地跨越样本边界进行注意力计算。

序列并行要求您的模型使用受支持的注意力实现之一（sdpa、flash_attention_2 或 flash_attention_3）。不支持 eager 注意力实现，因为它不能正确处理 position_ids。

使用序列并行时，请确保您的序列已正确填充。使用数据整理器中的 pad_to_multiple_of 来确保序列可以被 sp_size 整除。例如，对于 sp_size=4，请将 pad_to_multiple_of 设置为 4 或更高。

from transformers import DataCollatorForLanguageModeling

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=4,  # Ensure sequences are divisible by sp_size
)

当使用 sp_size 和多个 GPU 时，您**必须**显式设置 dp_replicate_size 或 dp_shard_size，以确保 total_size = dp_replicate_size * dp_shard_size * sp_size 等于您的总 GPU 数量。例如，对于 8 个 GPU 和 sp_size=4，您必须设置 dp_replicate_size=2（因为 2 × 1 × 4 = 8）。

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=4,
    dp_replicate_size=2,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
        sp_attn_implementation="flash_attention_2",
    ),
)

Trainer 会自动处理序列并行化的特殊要求，包括

  • 通过 DeepSpeed 的 UlyssesSPDataLoaderAdapter 适配数据加载器，以在 GPU 上分片序列。重要提示：与张量并行（所有 rank 必须接收相同的数据）不同，SP 中的每个 rank 从 DataLoader 接收唯一的数据流（类似于 DP）。适配器会在内部处理 SP rank 之间的序列块分发，因此您的 DataLoader 应继续向每个 rank 馈送不同的样本。
  • 生成 position_ids（如果未提供）
  • 为因果语言模型创建 shift_labels
  • 聚合跨序列并行 rank 的损失，并对 -100 标签进行适当的掩码

您可以使用 accelerate launch 命令启动带有序列并行化的训练。

accelerate launch --config_file alst_config.yaml your_training_script.py \
--output_dir output_dir \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 1

训练特性

DeepSpeed 支持许多可在配置文件中配置的训练特性。本节介绍其中一些最重要的特性。

梯度检查点

梯度检查点通过仅存储*一些*中间激活值而不是存储*所有*中间激活值来节省内存。这有助于在 GPU 上拟合更大的模型而不会耗尽内存，或者增加批大小以提高性能。但训练速度会变慢。

  • 对于 Transformers 模型，请设置 model.gradient_checkpointing_enable() 或在 TrainingArguments 中添加 --gradient_checkpointing。
  • 对于非 Transformers 模型，请使用 DeepSpeed 激活检查点 API。用 DeepSpeed API 替换 Transformers 的建模代码和 torch.utils.checkpoint 可以提供更大的灵活性，因为您可以将前向激活卸载到 CPU 内存，而不是重新计算它们。

批大小

批大小可以自动配置或手动设置。当您选择 "auto" 选项时，Trainer 会将 train_micro_batch_size_per_gpu 和 train_batch_size 设置为 world_size * per_device_train_batch_size * gradient_accumulation_steps 的值。

{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto"
}

通信数据类型

通信集体操作（如归约、收集和散布操作）使用单独的数据类型。

所有收集和散布操作都以数据存在的相同数据类型执行。例如，如果您正在使用 bf16 进行训练，那么数据也将以 bf16 收集，因为收集是非损耗操作。

归约操作是损耗性的，例如，当梯度跨多个 GPU 进行平均时。当通信以 fp16 或 bf16 进行时，它更有可能是损耗性的，因为低精度下多个数字的加法不精确。bf16 尤其如此，因为它的精度低于 fp16。因此，fp16 是归约操作的默认选项，因为平均梯度时的损耗很小。

通过在配置文件中设置 communication_data_type 参数来选择通信数据类型。例如，选择 fp32 会增加少量开销，但确保归约操作以 fp32 累积，并在准备就绪后，将其转换为您正在训练的任何半精度数据类型。

{
    "communication_data_type": "fp32"
}

梯度累积

梯度累积通过在多个 mini-batch 数据上累积梯度，然后再更新参数。它存储的梯度更少，并允许使用更大的*有效批大小*进行训练。训练速度会变慢，但有助于克服内存限制。

梯度累积可以自动配置或手动设置。当您选择 "auto" 选项时，Trainer 会将其设置为 gradient_accumulation_steps 的值。

{
    "gradient_accumulation_steps": "auto"
}

梯度裁剪

梯度裁剪有助于防止梯度爆炸，从而导致训练不稳定。它设置一个最大阈值，如果梯度的范数超过该阈值，则对梯度进行重新缩放。

梯度裁剪可以自动配置或手动设置。当您选择 "auto" 选项时，Trainer 会将其设置为 max_grad_norm 的值。

{
    "gradient_clipping": "auto"
}

混合精度训练

混合精度通过使用半精度进行一些计算来加速训练速度，但它也保留一些计算使用全精度来保持准确性。DeepSpeed 支持 fp32、fp16 和 bf16 数据类型。

fp32
fp16
bf16

如果模型未在混合精度下预训练，请使用 fp32 进行训练，因为它可能导致下溢或溢出错误。在这种情况下，请禁用 fp16（默认）。

{
    "fp16": {
        "enabled": false
    }
}

对于 Ampere GPU 和 PyTorch 1.7+，更高效的 tf32 模式会自动为某些操作启用，但结果仍为 fp32。通过在 Trainer 中设置 --tf32 来启用它，设置 --tf32 0 或 --no_tf32 来禁用它。

优化器和调度器

DeepSpeed 和 Transformers 的优化器和调度器可以混合搭配使用，前提是未启用 offload_optimizer。当启用 offload_optimizer 时，请使用非 DeepSpeed 的优化器（LAMB 除外），只要它同时具有 CPU 和 GPU 实现。

为配置文件设置优化器和调度器参数时，应从命令行进行，以避免难以查找的错误。例如，如果学习率在其他地方设置为不同的值，您可以通过命令行覆盖它。

优化器
scheduler

DeepSpeed 提供了几种 优化器（Adam、AdamW、OneBitAdam 和 LAMB），但您也可以从 PyTorch 导入其他优化器。如果您不在配置文件中配置优化器，Trainer 会自动选择 AdamW，并使用命令行中提供的或默认值来设置以下参数：lr、adam_beta1、adam_beta2、adam_epsilon、weight_decay。

您可以将参数设置为 "auto" 或手动输入自己的值。

{
   "optimizer": {
       "type": "AdamW",
       "params": {
         "lr": "auto",
         "betas": "auto",
         "eps": "auto",
         "weight_decay": "auto"
       }
   }
}

通过将以下内容添加到顶级配置中来使用不受支持的优化器。

{
   "zero_allow_untested_optimizer": true
}

从 DeepSpeed 0.8.3+ 开始，如果您想使用 offload，还需要将以下内容添加到顶级配置中，因为 offload 与 DeepSpeed 的 CPU Adam 优化器配合效果最佳。

{
   "zero_force_ds_cpu_optimizer": false
}

通用检查点

通用检查点可以在不同的模型架构、并行技术和训练配置之间保存和加载模型、优化器和训练调度器状态。通过将它们保存为通用格式，可以更轻松地继续模型训练和微调。

通过在配置文件中将 load_universal 设置为 true 来使用通用检查点恢复训练。

{
    "checkpoint": {
        "load_universal": true
    }
}

部署

DeepSpeed 可以通过其原生启动器、torchrun 或 Accelerate 进行部署。

在命令行中将 --deepspeed ds_config.json 参数添加到 Trainer。建议使用 DeepSpeeds add_config_arguments 工具将任何其他命令行参数添加到您的代码中。

多 GPU
单 GPU

要在多 GPU 上部署 DeepSpeed，请添加 --num_gpus。如果您打算使用所有可用的 GPU，则无需添加 --num_gpus。

deepspeed --num_gpus=2 examples/pytorch/translation/run_translation.py \
--deepspeed tests/deepspeed/ds_config_zero3.json \
--model_name_or_path google-t5/t5-small --per_device_train_batch_size 1 \
--output_dir output_dir --fp16 \
--do_train --max_train_samples 500 --num_train_epochs 1 \
--dataset_name wmt16 --dataset_config "ro-en" \
--source_lang en --target_lang ro

多节点

多节点设置由多个节点组成，每个节点有一个或多个 GPU 运行工作负载。DeepSpeed 需要共享存储系统，但如果不是这种情况，您需要调整配置文件以包含 检查点，以便在无法访问共享文件系统的情况下进行加载。

{
  "checkpoint": {
    "use_node_local_storage": true
  }
}

您还可以使用 TrainingArguments 中的 --save_on_each_node 参数将上述 checkpoint 自动添加到您的配置中。

下面 Torchrun 和 DeepSpeed 启动器的示例展示了如何部署两个节点，每个节点有八个 GPU。通过 ssh hostname1 访问第一个节点，通过 ssh hostname2 访问第二个节点。两个节点都必须能够通过 ssh 无密码本地通信。

torchrun
DeepSpeed

使用 torchrun，ssh 到每个节点并在两个节点上运行以下命令。启动器会等待两个节点同步后再启动训练。

torchrun --nproc_per_node=8 --nnode=2 --node_rank=0 --master_addr=hostname1 \
--master_port=9901 your_program.py <normal cl args> --deepspeed ds_config.json

Slurm

Slurm 是一个集群管理和作业调度系统。下面是一个 Slurm 脚本示例。

#SBATCH --job-name=test-nodes        # name
#SBATCH --nodes=2                    # nodes
#SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
#SBATCH --cpus-per-task=10           # number of cores per tasks
#SBATCH --gres=gpu:8                 # number of gpus
#SBATCH --time 20:00:00              # maximum execution time (HH:MM:SS)
#SBATCH --output=%x-%j.out           # output file name

export GPUS_PER_NODE=8
export MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
export MASTER_PORT=9901

srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
 --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
 --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
your_program.py <normal cl args> --deepspeed ds_config.json'

使用以下命令同时在所有节点上启动训练。

sbatch launch.slurm

Jupyter Notebook

要在 Jupyter Notebook 中使用 DeepSpeed，您需要模拟一个分布式环境，因为启动器不支持从笔记本启动。这仅支持单 GPU。要使用多 GPU，您必须使用多进程环境，这意味着您必须使用 DeepSpeed 启动器，而它无法像此处所示那样进行模拟。

# emulate a launcher in the notebook
import os

os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = "9994"  # modify if RuntimeError: Address already in use
os.environ["RANK"] = "0"
os.environ["LOCAL_RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"

training_args = TrainingArguments(..., deepspeed="ds_config_zero3.json")
trainer = Trainer(...)
trainer.train()

在笔记本中，在当前目录创建一个临时配置文件，使用一个单独的单元格。

%%bash
cat <<'EOT' > ds_config_zero3.json
{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },

    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": "auto",
            "betas": "auto",
            "eps": "auto",
            "weight_decay": "auto"
        }
    },

    "scheduler": {
        "type": "WarmupLR",
        "params": {
            "warmup_min_lr": "auto",
            "warmup_max_lr": "auto",
            "warmup_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,
        "sub_group_size": 1e9,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_max_live_parameters": 1e9,
        "stage3_max_reuse_distance": 1e9,
        "stage3_gather_16bit_weights_on_model_save": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
}
EOT

如果训练脚本在一个文件中而不是一个笔记本单元格中，请在笔记本单元格中使用 shell 启动 DeepSpeed。

!git clone https://github.com/huggingface/transformers
!cd transformers; deepspeed examples/pytorch/translation/run_translation.py ...

另一种选择是使用 %%bash 来运行 shell 程序，而无需模拟分布式环境。但是，在训练完成之前，您将无法看到日志。

%%bash

git clone https://github.com/huggingface/transformers
cd transformers
deepspeed examples/pytorch/translation/run_translation.py ...

保存模型权重

DeepSpeed 将主要的 fp32 权重存储在自定义检查点优化器文件中（global_step*/*optim_states.pt），这些文件保存在普通检查点下。

fp16

ZeRO-2 将模型权重保存在 fp16 中。要为 ZeRO-3 将权重保存在 fp16 中，请在配置文件中设置 "stage3_gather_16bit_weights_on_model_save": true，因为权重分布在多个 GPU 上。

如果不这样做，Trainer 将不会以 fp16 格式保存权重，也不会创建 pytorch_model.bin 文件。这是因为 DeepSpeed 的 state_dict 包含一个占位符而不是实际权重，因此您将无法加载它。

{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}

fp32

除非您有大量的可用 CPU 内存，否则不应在训练期间保存 fp32 权重，因为它可能需要大量内存。通常最好在训练完成后离线保存 fp32 权重。

离线
online

DeepSpeed 在顶级检查点文件夹中提供了一个 zero_to_fp32.py 脚本，用于在任何时候提取权重。这是一个独立的脚本，您不需要配置文件或 Trainer。

例如，如果您的检查点文件夹如下所示，那么您可以通过运行以下命令从多个 GPU 创建并合并 fp32 权重到一个 pytorch_model.bin 文件中。该脚本会自动发现包含检查点的子文件夹 global_step1。

$ ls -l output_dir/checkpoint-1/
-rw-rw-r-- 1 stas stas 1.4K Mar 27 20:42 config.json
drwxrwxr-x 2 stas stas 4.0K Mar 25 19:52 global_step1/
-rw-rw-r-- 1 stas stas   12 Mar 27 13:16 latest
-rw-rw-r-- 1 stas stas 827K Mar 27 20:42 optimizer.pt
-rw-rw-r-- 1 stas stas 231M Mar 27 20:42 pytorch_model.bin
-rw-rw-r-- 1 stas stas  623 Mar 27 20:42 scheduler.pt
-rw-rw-r-- 1 stas stas 1.8K Mar 27 20:42 special_tokens_map.json
-rw-rw-r-- 1 stas stas 774K Mar 27 20:42 spiece.model
-rw-rw-r-- 1 stas stas 1.9K Mar 27 20:42 tokenizer_config.json
-rw-rw-r-- 1 stas stas  339 Mar 27 20:42 trainer_state.json
-rw-rw-r-- 1 stas stas 2.3K Mar 27 20:42 training_args.bin
-rwxrw-r-- 1 stas stas 5.5K Mar 27 13:16 zero_to_fp32.py*

运行 python zero_to_fp32.py -h 可获得更多用法详情。该脚本需要 2 倍于最终 fp32 权重的总 RAM。

python zero_to_fp32.py . pytorch_model.bin

非 Trainer 集成

DeepSpeed 也可以与 Transformers 一起使用，而无需 Trainer。在调用 from_pretrained() 时，HfDeepSpeedConfig 负责收集 ZeRO-3 参数并将模型分区到多个 GPU 上。

为了有效地部署 ZeRO-3，您必须在加载模型之前实例化 HfDeepSpeedConfig。

预训练模型
非预训练模型
from transformers.integrations import HfDeepSpeedConfig
from transformers import AutoModel
import deepspeed

# DeepSpeed config object or path to the file
ds_config = {...}
# must run before instantiating the model to detect ZeRO-3
dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive
model = AutoModel.from_pretrained("openai-community/gpt2")
engine = deepspeed.initialize(model=model, config_params=ds_config, ...)

故障排除

当遇到错误时，首先要做的事情之一是检查 DeepSpeed 是否是原因（因为很多时候不是）。尝试在没有 DeepSpeed 的情况下重试您的设置，如果错误仍然存在，请报告该问题。如果问题与 Transformers 集成无关，请在 DeepSpeed 仓库中打开问题。

对于与 Transformers 集成相关的问题，请提供以下信息。

  • 完整的 DeepSpeed 配置文件。

  • 如果您正在自己编写脚本设置 Trainer，请提供 Trainer 的命令行参数或 TrainingArguments（不要转储包含许多无关条目的整个 TrainingArguments）。

  • 以下命令的输出。

    python -c 'import torch; print(f"torch: {torch.__version__}")'
python -c 'import transformers; print(f"transformers: {transformers.__version__}")'
python -c 'import deepspeed; print(f"deepspeed: {deepspeed.__version__}")'

  • 一个可用于重现问题的 Google Colab 笔记本链接。

  • 一个标准或非自定义数据集，或者一个现有的示例，用于重现问题。

以下部分将指导您解决两个最常见的问题。

启动时进程被杀死

当 DeepSpeed 进程在启动时被杀死而没有堆栈跟踪时，这通常意味着程序试图分配比您的系统上可用的 CPU 内存更多的内存。或者，进程可能试图分配比允许的更多的 CPU 内存，导致操作系统内核终止该进程。

在这种情况下，请检查您的配置文件是否配置了 offload_optimizer、offlload_param 或两者都配置为卸载到 CPU。

如果您设置了 NVM3 和 ZeRO-3，请尝试卸载到 NVMe（首先 估算模型的内存需求），而不是。

NaN 损失

NaN 损失通常发生在模型以 bf16 预训练，而您尝试使用 fp16（尤其对于 TPU 训练的模型）时。要解决此问题，请使用 fp32 或 bf16（如果您的硬件支持，例如 TPU、Ampere GPU 或更新的 GPU）。

fp16 也有可能导致溢出。例如，如果您的配置文件如下所示，您可能会在日志中看到以下溢出错误。

{
    "fp16": {
        "enabled": "auto",
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    }
}

下面的 OVERFLOW! 错误是 DeepSpeed 损失缩放器无法找到缩放系数来克服损失溢出造成的。在这种情况下，尝试一个更高的 initial_scale_power 值（32 通常有效）。

0%|                                                                                                                             | 0/189 [00:00<?, ?it/s]
 [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 262144, reducing to 262144
  1%|▌                                                                                                                    | 1/189 [00:00<01:26,  2.17it/s]
 [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 262144, reducing to 131072.0
  1%|█▏
 [...]
 [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1
 14%|████████████████▌                                                                                                   | 27/189 [00:14<01:13,  2.21it/s]
 [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1
 15%|█████████████████▏                                                                                                  | 28/189 [00:14<01:13,  2.18it/s]
 [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1
 15%|█████████████████▊                                                                                                  | 29/189 [00:15<01:13,  2.18it/s]
 [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 1, reducing to 1
[...]

资源

DeepSpeed 是用于大规模模型训练的强大技术。要了解更多关于 DeepSpeed 的信息，请查看他们的 博客文章、文档和 GitHub。

以下论文提供了关于 ZeRO 的更多详细信息。

在 GitHub 上更新

←FullyShardedDataParallel 多 GPU 调试→
© . This site is unofficial and not affiliated with Hugging Face, Inc.