Colossal-AI发布新的深度学习辅助框架，半小时即可完成ViT-Base/32预训练

从头开始预训练计算机视觉领域最热门的 AI 模型 Vision Transformer (ViT) 需要多长时间？ Colossal-AI系统最新给出的答案是半小时！

近年来，随着AI性能的不断提升，AI模型参数的数量也呈现爆发式增长，从AlexNet、ResNet到BERT、GPT、MoE……，AI模型的参数量级不断刷新，现已超过万亿 ，这使得训练成本急剧上升。 仅仅训练2020年发布的千亿参数GPT-3，即便是用一块最先进的英伟达A100 GPU进行训练，也需要等待100多年。 显然，如何加速AI训练过程已经成为AI行业最大的痛点之一。

加速AI模型训练的一种常见方法是借助数据并行性来实现大批量训练，即通过增加批量大小，从而减少迭代次数，最终实现训练时间的显着减少。 但是由于泛化差距，大批量训练通常会导致测试集上的收敛困难和性能下降。

因此，越来越多的科技巨头选择使用高性能计算技术，借助谷歌的 TPU Pod 和英伟达的 SuperPOD 等超级计算机集群，使用数百甚至数千个最好的处理器来并行加速计算任务。

然而，即使是价值数亿美元的超级计算机，也面临着当硬件堆栈达到一定数量时，计算效率无法进一步提高，浪费大量计算资源的瓶颈。 此外，分布式并行编程通常需要与计算机系统和架构相关的专业知识，这进一步增加了前沿人工智能模型的训练成本。

在 ImageNet-1K 数据集（128 万张图像）上，ViT 的预训练需要使用整个数据集 300 次，使用一个批处理大小为 128 的 NVIDIA A100 GPU 预训练 ViT-Base/32 大约需要三天时间 .

使用 Colossal-AI 提供的大批量优化器 LAMB，可以成功克服大批量优化的困难，通过 32K 批量大小的数据并行快速将训练过程扩展到 200 个 A100 GPU，只需 0.61 小时即可完成 ViT-Base/32 训练并保持准确性。 对于更复杂的ViT-Base/16 和 ViT-Large/32，Colossal-AI 也分别只需 1.87 小时和 1.30 小时即可完成。

目前主流的张量并行方案，如微软的DeepSpeed和英伟达的Megatron，都是一维张量并行，这意味着每个处理器都需要与其他所有处理器进行通信。但是，处理器之间的通信，尤其是多机跨节点通信，不仅会严重降低运行速度，占用处理器大量额外内存，而且能耗成本非常高。针对这个痛点，Colossal-AI 提出使用多维张量并行：2/2.5/3 维的张量并行。 Colossal-AI还完全兼容现有的数据并行、流水线并行、序列并行等并行模式，共同构成了Colossal-AI的“六维并行”，远超微软和英伟达的三维并行，相关技术可以轻松扩展到推理加速领域。

Colossal-AI作为深度学习框架的核心，将“系统优化”与“上层框架”和“下层硬件”解耦，为多维并行、大批量优化器、零等前沿技术的实现提供了便捷的实现。 以即插即用的方式卸载冗余内存。 易于扩展和使用，只需要少量的代码修改，防止用户学习复杂的分布式系统知识。 对于GPU等处理器来说，使用密集集群的平均算力成本会比分散使用低很多，这样不仅可以节省训练时间，还可以大大降低算力成本。

Colossal-AI 仍处于 beta 测试阶段，在不久的将来会发布更多惊人的结果。

Paper Address: https://arxiv.org/abs/2110.14883

Project Address: https://github.com/hpcaitech/ColossalAI

Document Address: https://www.colossalai.org/

作者:HPC-AI Tech