Huggingface - 4.DeepSpeed

对于 deepspeed 这个库，不打算直接单独使用这个库，而是将其结合 transformers 或者 accelerate 一起使用。

该文档就是介绍如何使用集成到 transformers 库中的 deepspeed。

1.Trainer

在 transformers 的 Trainer 库中集成了 deepspeed，上层应用只要使用了 Trainer 之后基本上是什么都不用管，直接给定 deepspeed 的配置文件即可，非常方便。

推荐使用这种方法！

推荐使用这种方法！

2.配置文件

配置文件这一块，有三点：

• 本文中会给出一些配置文件的例子；
• 如何在项目 DeepSpeedExamples 中找相应的配置文件的例子；
• 配置文件中各个参数的含义与作用，这个主要见 DeepSpeed 的官方文档：DeepSpeed Configuration JSON

2.1 DeepSpeedExamples

项目 DeepSpeedExamples 中给了不少的使用 deepspeed 的配置文件，可以在该项目中查找符合自己需要的配置文件。这个操作非常简单，就是使用 find 命令在这个repo中寻找相应的配置文件，命令如下：

git clone https://github.com/microsoft/DeepSpeedExamples
cd DeepSpeedExamples
find . -name '*json'

另外还可以结合上 grep 命令，比如查找关键字 Lamb 的样例如下：

grep -i Lamb $(find . -name '*json')

2.2 配置文件样例

配置 ZeRO stage 为 2，开启 cpu offload 功能，使用 AdamW 优化器和 WarmupLR scheduler，并且使用混合精度训练的配置文件如下所示：

{
    "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": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 2e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },

    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
}

3.共享配置

这部分其实是一个注意点，在 transformers 项目中的 Trainer 会接收很多个参数，DeepSpeed 也会接收很多个参数，而且这两个不同的工具在指定同一个参数时，参数名可能还会不一样。这样就有很大的可能出现同一个参数在 DeepSpeed 中配置的值和 Trainer 中配置的值是不同的，如果出现这种情况，可能会发生一些未知错误，并且很难排查。

在这篇 huggingface 的文档中给出了避免上述问题的解决方案是：在 deepspeed 的配置文件中，可以将部分参数配置为 auto，见上一小节中举的配置文件的例子。当 deepspeed 的配置文件这些参数配置为了 auto 时，那么实际生效的就是 Trainer 中这些参数设置的值，这样就避免了同一个参数给设定了不同值的问题了。

4.ZeRO（Zero Redundancy Optimizer）

先说一下 ZeRO 这个工具几个不同状态的功能：

• ZeRO stage 0：不使用ZeRO；
• ZeRO stage 1：对优化器进行分片；
• ZeRO stage 2：对优化器进行分片+对梯度进行分片；
• ZeRO stage 3：对优化器进行分片+对梯度进行分片+对模型参数进行分片；

其中比较常用的是 ZeRO2 和 ZeRO3。

4.2 ZeRO2 config

该部分参考的文档为：* https://huggingface.co/docs/transformers/main_classes/deepspeed#zero2-config

配置样例如下：

{
    "zero_optimization": {
        "stage": 2,
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": true,
        "reduce_scatter": true,
        "reduce_bucket_size": 5e8,
        "contiguous_gradients": true
    }
}

其中 allgather_bucket_size、overlap_comm、reduce_bucket_size 这几个参数的作用没有看，如果需要使用可以阅读该部分的参考文档，这几个参数都是用来优化性能的。

4.3 ZeRO3 config

该部分参考的文档为：* https://huggingface.co/docs/transformers/main_classes/deepspeed#zero3-config

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

同样的，其中的参数 stage3_prefetch_bucket_size、stage3_param_persistence_threshold、stage3_max_live_parameters、stage3_max_reuse_distance 都没有看，都是优化性能的参数，现在的阶段是先跑起来，之后再来看这些参数的作用。

4.4 ZeRO0 和 ZeRO1

应该不怎么使用这两个配置，略。

5.Optimizer and Scheduler#

该部分参考的文档为：* https://huggingface.co/docs/transformers/main_classes/deepspeed#optimizer-and-scheduler

关于 optimizer 和 scheduler 部分，这里还有一个能否混用的问题，就是 deepspeed 库中实现了部分的 optimizer 和 scheduler，然后 transformers 库中也实现了部分的 optimizer 和 scheduler，那么他们能不能混合使用？能否混合使用的具体情况如下表所示，但是这里感觉上没必要混着使用，全部使用 deepspeed 的就可以了。

Combos	HF Scheduler	DS Scheduler
HF Optimizer	Yes	Yes
DS Optimizer	No	Yes

5.1 Optimizer

deepspeed 支持的优化器有 Adam、AdamW、OneBitAdam、Lamb 这几个。当然了，一般情况下，目前主要使用的优化器还是 Adam。

下面是优化器 Adam 的配置：

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

除了将各个参数配置为 auto 以外，还可以指定具体的值，如下所示：

{
   "optimizer": {
       "type": "AdamW",
       "params": {
         "lr": 0.001,
         "betas": [0.8, 0.999],
         "eps": 1e-8,
         "weight_decay": 3e-7
       }
   }
}

5.2 Scheduler

deepspeed 支持的 scheduler 有 LRRangeTest、OneCycle、WarmupLR、WarmupDecayLR 这几个。当然了，一般情况下，目前主要使用的 scheduler 还是 WarmupLR。

配置举例如下：

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

如果指定具体的值的话，配置如下：

{
   "scheduler": {
         "type": "WarmupLR",
         "params": {
             "warmup_min_lr": 0,
             "warmup_max_lr": 0.001,
             "warmup_num_steps": 1000
         }
     }
}

Reference#

• https://huggingface.co/docs/transformers/main_classes/deepspeed
• https://huggingface.co/docs/accelerate/usage_guides/deepspeed