IBM联合PyTorch大幅改进在云计算训练大型AI模型的效率

IBM研究院与PyTorch合作，发展一个称作rate_limiter的控件，能够通过配置训练时的内存，使得具有数十亿参数的模型，也能够在标准云计算网络基础设施，像是以太网络上运行。

当前的人工智能模型，只要使用简单的文本题事，就能创造出图像、歌曲甚至是网站，而这种具有数十亿参数的模型，被称为基础模型，不需要大量时长的训练与标记，基础模型就可以被简单地重新用于另一项任务。

基础模型主要在高端的高性能计算（HPC）基础设施上运行，虽然这些系统很可靠，但IBM提到，对于许多想要自己训练基础模型，并且仅供自己使用的人来说，HPC设备是一个极高的门槛，包括高端GPU搭配低延迟InfiniBand网络系统，以及各种专门的操作流程等。

而IBM研究人员与PyTorch分布式团队合作，寻找出可在网络硬件上，训练大型人工智能模型的方法，而这个联合小组也已经验证，可以在Red Hat的OpenShift平台上，使用基于以太网络的常规网络，来扩展和训练大型模型。

使用PyTorch的完全分片资料平行技术（FSDP），该团队可以在IBM Cloud上的标准以太网络，训练具有110亿参数的模型，实现与HPC网络系统同等规模的高效率训练。研究人员提到，过去尝试使用以太网络，在PyTorch上训练数十亿参数的模型，但是效果不佳，远低于训练基础模型所需要的性能，随着模型规模增长，只有当GPU能够完整包存模型副本和训练状态时，资料平行训练的标准方法才能发挥作用。

虽然FSDP或DeepSpeed等新方法可以在训练期间，有效地将模型和资料分散到多个GPU上，但也只限于HPC系统，而无法在由以太网络连接的系统上运行。为此，联合研究团队开发了一个FSDP API，并构建rate_limiter控件，该组件可以控制用于发送和接收张量所使用的内存，进而缓解系统的内存压力，使效率较之前提高达4.5倍。

这项研究使用的基础设施为现成的硬件，该系统在IBM Cloud上运行，由200个节点组成，每个节点有8张Nvidia A100 80GB显卡，与96 vCPU、1.2TB CPU内存，节点内的显卡以NVLink连接，显卡之间带宽为600 GBps，而节点之间通过两条100 Gbps以太网络连接，提供120 GBps可用带宽。

该GPU系统自5月来持续运行，研究人员配置Red Hat OpenShift容器平台执行人工智能负载，用于大型人工智能模型的端到端训练、微调和推理。研究人员提到，这是目前该产业中，第一个使用Kubernetes、PyTorch FSDP API和标准以太网络，高效率实现高达110亿个参数模型的用例，而这将有助于其他组织以更具成本效益的方式，在云计算训练大型模型。