Hugging Face 的 LOGO

Transformers 文档

DeepSpeed

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

DeepSpeed

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

ZeRO 阶段 描述
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 配置 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 文件的路径传递；如果您在 Jupyter Notebook 中使用 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 配置 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
        },
        "aio": {
            "block_size": 262144,
            "queue_depth": 32,
            "thread_count": 1,
            "single_submit": false,
            "overlap_events": 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
}

训练功能

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

梯度累积

梯度累积在更新参数之前，会在多个小批次数据上累积梯度。它存储的梯度较少，并允许使用更大的*有效批次大小*进行训练。不过，训练速度会变慢，但对于克服内存限制很有用。

梯度累积可以自动配置或手动设置。当您选择 `“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` 禁用。

优化器和调度器

如果未启用 offload_optimizer，DeepSpeed 和 Transformers 优化器和调度器可以混搭。当启用 offload_optimizer 时，只要它具有 CPU 和 GPU 实现，就可以使用非 DeepSpeed 优化器（LAMB 除外）。

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

优化器
调度器

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+ 开始，如果您想使用卸载，还需要将以下内容添加到顶级配置中，因为卸载与 DeepSpeed 的 CPU Adam 优化器配合得最好。

{
   "zero_force_ds_cpu_optimizer": false
}

通用检查点

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

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

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

部署

DeepSpeed 可以使用其原生启动器 torchrun 或 Accelerate 进行部署。

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

多 GPU
单 GPU

要将 DeepSpeed 部署到多个 GPU 上，请添加 `--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 --overwrite_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，您需要模拟分布式环境，因为启动器不支持从 Notebook 部署。这仅支持单个 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()

在 Notebook 的当前目录中，通过专用单元格即时创建配置文件。

%%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

如果训练脚本在文件中而不是在 notebook 单元格中，则从 notebook 单元格中的 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，如果模型权重分布在多个 GPU 上，请在配置文件中设置 `"stage3_gather_16bit_weights_on_model_save": true` 以 fp16 格式保存权重。

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

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

fp32

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

离线
在线

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` 以获取更多用法详情。该脚本需要最终 fp32 权重两倍的通用 RAM。

python zero_to_fp32.py . pytorch_model.bin

非 Trainer 集成

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

您必须在加载模型之前实例化 HfDeepSpeedConfig，以高效部署 ZeRO-3。

预训练模型
非预训练模型
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 或 TrainingArguments（如果您自己编写 Trainer 设置脚本）（请勿转储包含许多不相关条目的整个 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`、`offload_param` 或两者都卸载到 CPU。

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

NaN 损失

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

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 调试→