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 还提供了成本监控和分析等能力，能够面向任务去管理和优化分布式训练的成本。