关于阿里云:DeepSpeed-Kubernetes-如何轻松落地大规模分布式训练

作者：刘霖

背景现状

随着 ChatGPT 的广泛应用，各种大规模语言模型层出不穷，其中包含 EleutherAI 推出的 200 亿参数的 GPT-NeoX-20B 和 BigScience 公布的 1760 亿参数的 Bloom 模型。

因为模型越来越大，单张 GPU 已无奈加载整个模型，分布式模型训练成为了一种必然的趋势。在 GPT-NeoX 和 Bloom 的背地，DeepSpeed 框架是实现分布式模型训练的要害。

DeepSpeed 是一个开源的深度学习训练优化库，提供了多种优化策略，如混合精度训练、数据并行、模型并行、流水线并行等，这些策略可用于减速大规模模型的训练。此外，DeepSpeed 还提供了高性能的分布式训练框架，反对多种支流的深度学习框架，并且能够在不同的硬件和云平台上进行训练。借助 DeepSpeed，算法工程师能够更加疾速地训练大规模深度学习模型，从而进步模型的准确性和效率。

以后，越来越多企业在云上基于容器和 Kubernetes 进行大规模的分布式深度学习训练，充分利用弹性、扩展性、自动化、高可用等劣势，大幅提高分布式训练的效率和可靠性，同时升高治理老本和复杂性。

然而，随着模型规模扩充以及企业对生产效率的一直谋求，将 DeepSpeed 分布式训练任务在 Kubernetes 中搭建和运行依然存在着很多挑战和难点。例如，GPU 资源利用率低，分布式训练扩展性差，以及难以不便地获取实时日志和监控等。

计划介绍

目前，阿里云容器服务 ACK 云原生 AI 套件曾经反对 DeepSpeed 分布式训练，为您提供高效便捷的解决方案。

您只需筹备好训练代码和数据，就能够利用命令行工具 Arena 疾速在 ACK 集群中部署基于 DeepSpeed 的分布式训练任务。此外，能够通过 TensorBoard 可视化工具不便地查看训练作业的状态和后果，从而使 DeepSpeed 分布式训练变得更加容易和高效。

对于 ACK 云原生 AI 套件的更多信息，请通过往期文章进行理解： 解脱 AI 生产“小作坊”：如何基于 Kubernetes 构建云原生 AI 平台

外围劣势

基于阿里云容器服务 ACK 云原生 AI 套件搭建和运行 DeepSpeed 分布式训练任务具备以下劣势：

1. 大规模异构资源管理

应用容器服务 ACK，您能够轻松治理大规模异构资源，疾速搭建基于 CPU、GPU、FPGA 等不同类型计算资源的规范 Kubernetes 集群，实现对各类异构资源进行对立、灵便的治理、调度和运维。云原生 AI 套件还反对多种 GPU 调度策略（共享 + 隔离、优先级、拓扑感知等）和丰盛的 GPU 监控告警等能力，帮您进一步优化资源利用率。

2. 灵便弹性与老本优化

应用 ACK 弹性节点池和 HPA/VPA，您能够灵便按需主动伸缩 GPU 节点数和 Pod 数，并基于 GPU 指标实现 AHPA 弹性预测。云原生 AI 套件反对 ECS、ECI 等弹性资源混合的高级调度，包含 affinity/anti-affinity、pod topology spreading、deploymentset 感知等调度。此外，通过反对资源有余到时终止、checkpoint 主动保留、容错、Failover 等能力，解决基于抢占式实例进行分布式训练的可用性问题，缩小训练老本的同时，近乎不影响训练作业成功率。ACK 还提供了老本监控和剖析等能力，可能面向工作去治理和优化分布式训练的老本。

3. 高效的工作治理调度

云原生 AI 套件提供命令行工具 Arena，对深度学习外围生产环节（包含数据管理、模型开发、模型训练、模型评估、推理服务部署等）工作进行简略形象和高效治理。通过 Arena，您能够疾速提交分布式训练任务，并晋升训练任务从提交到运行的性能，以及进行工作的生命周期治理。此外，云原生 AI 套件还提供针对分布式场景优化的调度策略，例如按 GPU 卡的 Binpack 算法调配策略，晋升 GPU 卡利用率，还反对自定义的工作优先级治理和租户弹性资源配额管制，在确保用户资源分配的根底上，通过资源共享的形式来晋升集群的整体资源利用率。

疾速应用

上面将介绍如何基于阿里云容器服务 ACK 云原生 AI 套件疾速搭建和运行 DeepSpeed 分布式训练：

前提条件

• 已创立蕴含 GPU 的 Kubernetes 集群。具体操作，请参见创立蕴含 GPU 的 Kubernetes 集群 [ 1]。
• 已装置云原生 AI 套件（ack-arena 版本不低于 0.9.6）。具体操作，请参见部署云原生 AI 套件 [ 2]。
• 已装置 Arena 客户端（版本不低于 0.9.6）。具体操作，请参见装置 Arena [ 3]。
• 已给集群配置了 Arena 应用的 PVC，具体操作，请参见配置 NAS 共享存储 [ 4]（或者配置 CPFS 共享存储 [ 5]）。

应用阐明

本示例应用 DeepSpeed 训练一个掩码语言模型（Masked Language Model）。为不便运行，已将示例代码 [ 6] 和数据集下载至示例镜像中；若您无需应用示例镜像，也反对从 Git URL 下载源代码，并将数据集寄存在共享存储系统（基于 NAS 的 PV 和 PVC）中。本示例假如您曾经取得了一个名称为 training-data 的 PVC 实例（一个共享存储），用来寄存训练后果。

如需自定义训练镜像，请参考如下形式：

• 形式一：

参考该 Dockerfile [ 7]，在根底镜像中装置 OpenSSH

• 形式二：

应用 ACK 提供的 DeepSpeed 根底镜像：

registry.cn-beijing.aliyuncs.com/acs/deepspeed:v072_base

示例介绍

本示例中创立了一个 Transformer Bert 模型，依据句子前后的上下文来填充句子。比方在下列句子中：

In the beautiful season of ____ the ____ shed their leaves.

依据提供的 ’In the beautiful season of’ 和 ‘shed their leaves’，能够预测到空白处应该填入 ‘Autumn’ 和 ‘trees’。

本示例通过整合 DeepSpeed 的能力来进步训练的速度和效率，在下列几个方面进行了优化：

• 混合精度训练：DeepSpeed 反对应用 fp16 数据类型进行混合精度训练。通过在 ds_config 中设置以下配置，即可启用混合精度训练。

"fp16": {"enabled": True}

• ZeRO 数据并行：Zero Redundancy Optimizer（零冗余优化器）能够反对每个 GPU 都只存储模型参数、梯度和优化器状态的一部分，从而升高 GPU 显存占用，反对更大的模型。以后反对 3 个阶段，阶段 1 对优化器状态进行分片。阶段 2 还会对梯度进行分片。阶段 3 进一步对模型权重进行分片。通过在 ds_config 中设置以下配置，即可启动阶段 1。

"zero_optimization": {"stage": 1}

• ZeRO-Offload：通过同时利用 GPU 和 CPU 的计算和存储资源，比方将优化器状态和梯度保留在内存上，从而使单 GPU 能够反对的模型更大。比方在一张 P40 GPU 上，无奈训练一个 20 亿参数的模型，然而应用 ZeRO-Offload 能够做到。通过在 ds_config 中设置以下配置，即可启用 ZeRO-Offload。

"zero_optimization": {
  "offload_optimizer": {"device": "cpu"}
}

本示例中 DeepSpeed 的残缺配置文件 ds_config 参考如下。

ds_config = {
    "train_micro_batch_size_per_gpu": batch_size,
    "optimizer": {
        "type": "Adam",
        "params": {"lr": 1e-4}
    },
    "fp16": {"enabled": True},
    "zero_optimization": {
        "stage": 1,
        "offload_optimizer": {"device": "cpu"}
    }
}

DeepSpeed 是基于 AllReduce 的分布式训练框架，在 hostfile 中保留了所有的 worker 信息，luancher 读取 hostfile 文件获取 worker 信息，通过 PDSH 的形式在每个 worker 上启动训练任务。

ACK 云原生 AI 套件中的 Operator 组件会主动生成上述配置。筹备好模型训练代码和数据后，即可通过 Arena 提交 DeepSpeed 分布式训练任务。如下图所示，Operator 会创立一个 Launcher Pod 和多个 Worker Pod 并在 Worker Pod 上启动训练任务。

操作步骤

1. 提交 DeepSpeed 作业

通过以下代码示例提交蕴含 1 个 Launcher 节点，3 个 Worker 节点的 DeepSpeed 训练任务。该工作将应用 3 台机器，每台机器上的一张 GPU 卡进行训练。

arena submit etjob \
    --name=deepspeed-helloworld \
    --gpus=1 \
    --workers=3 \
    --image=registry.cn-beijing.aliyuncs.com/acs/deepspeed:hello-deepspeed \
    --data=training-data:/data \
    --tensorboard \
    --logdir=/data/deepspeed_data \
    "deepspeed /workspace/DeepSpeedExamples/HelloDeepSpeed/train_bert_ds.py --checkpoint_dir /data/deepspeed_data"

预期输入如下：

trainingjob.kai.alibabacloud.com/deepspeed-helloworld created
INFO[0007] The Job deepspeed-helloworld has been submitted successfully
INFO[0007] You can run `arena get deepspeed-helloworld --type etjob` to check the job status

2. 执行以下命令获取作业详情

arena get deepspeed-helloworld

预期输入如下：

Name:      deepspeed-helloworld
Status:    RUNNING
Namespace: default
Priority:  N/A
Trainer:   ETJOB
Duration:  6m

Instances:
  NAME                           STATUS   AGE  IS_CHIEF  GPU(Requested)  NODE
  ----                           ------   ---  --------  --------------  ----
  deepspeed-helloworld-launcher  Running  6m   true      0               cn-beijing.192.1xx.x.x
  deepspeed-helloworld-worker-0  Running  6m   false     1               cn-beijing.192.1xx.x.x
  deepspeed-helloworld-worker-1  Running  6m   false     1               cn-beijing.192.1xx.x.x
  deepspeed-helloworld-worker-2  Running  6m   false     1               cn-beijing.192.1xx.x.x

Your tensorboard will be available on:
http://192.1xx.x.xx:31870

3. 通过浏览器查看 TensorBoard

执行以下命令，将集群中的 TensorBoard 映射到本地 9090 端口。

kubectl port-forward svc/deepspeed-helloworld-tensorboard 9090:6006

在浏览器中拜访 localhost:9090，即可查看 TensorBoard。

执行以下命令可获取作业日志信息：

arena logs deepspeed-helloworld

更多操作命令及参数解释可参考云原生 AI 套件文档 [ 8]。

相干链接：

[1] 创立蕴含 GPU 的 Kubernetes 集群

https://help.aliyun.com/document_detail/171074.html?spm=a2c4g…

[2] 部署云原生 AI 套件

https://help.aliyun.com/document_detail/201997.htm#section-2a…

[3] 装置 Arena

https://help.aliyun.com/document_detail/212117.htm#task-1917487

[4] 配置 NAS 共享存储

https://help.aliyun.com/document_detail/212177.htm#task-1917776

[5] 配置 CPFS 共享存储

https://help.aliyun.com/document_detail/212178.htm#task-1917776

[6] 示例代码

https://github.com/microsoft/DeepSpeedExamples/tree/master/training/HelloDeepSpeed

[7] Dockerfile

https://github.com/kubeflow/arena/blob/master/samples/deepspeed/Dockerfile

[8] 云原生 AI 套件文档

https://help.aliyun.com/document_detail/2249322.html?spm=a2c4…

如果您心愿深刻理解更多对于 ACK 云原生 AI 套件的信息，或者须要与咱们就 LLM/AIGC 等相干需要进行交换，欢送退出咱们的钉钉群：33214567。