IBM研发开源云原生AI训练网络

近年来，人工智能（AI）领域发展迅速，得益于人工智能模型不断成熟。这些模型已经可以用简单的文字提示制作出图像、歌曲甚至网站。这类具有数十亿个参数的模型被称为基础模型，只需很少的微调，就可以从一个任务转换到另一个任务，以承担新的应用需求。省去了无数小时的训练和标记。

基础模型主要是在高性能计算(HPC)基础设施上训练的，虽然可靠，但对于许多希望训练基础模型以供自己使用的人来说，这是一个昂贵的进入门槛。这些用于训练AI模型的系统必须是定制设计的，很少依赖于商用硬件选项。顶级的GPU与低延迟的InfiniBand（无限带宽技术，是一个用于高性能计算的计算机网络通信标准，它具有极高的吞吐量和极低的延迟）网络系统配套使用，而后者的设置和运行成本很高，还需要定制操作流程，从而进一步提高了成本。

解开高性能计算枷锁

IBM的研究人员一直在与PyTorch(Linux基金会运行的开源机器学习平台)中的分布式团队合作，寻找一种在负担得起的网络硬件上训练大型AI模型的方法。该小组的研究表明，在红帽公司的OpenShift平台上使用基于常规以太网的网络，可以扩展和训练大型模型。

通过PyTorch的完全分片数据并行（FSDP）技术，该团队能够使用IBM Cloud上的标准以太网网络，成功地训练具有110亿个参数的模型。IBM的方法使得以往只能运行中小型AI模型的HPC网络系统，实现了同等规模的高效训练模型。

分配内存以提高性能

以前尝试在PyTorch上用Ethernet训练具有数十亿个参数的模型，结果性能很差，远远低于训练基础模型所需的性能。在云计算中，期望在任何时候系统都被完全分配。随着AI模型越来越大，数据并行训练的标准方法只有在GPU能够保存模型的完整副本及其训练状态时才能工作。虽然新的训练技术——如PyTorch的FSDP或DeepSpeed——可以在训练过程中有效地将模型和数据分布到多个GPU上，但它们只能在高性能计算系统(HPC)上有效地运行，而不能在以太网连接的系统上有效地运行。联合团队探索了FSDP的API，并构建了一个名为rate_limititer的新控件，该控件控制了用于发送和接收张量的内存分配，减轻了系统的内存压力，并将缩放效率提高了4.5倍。

民主化访问云网络

团队用于这项工作的基础设施基本上是现成的硬件。该系统运行在IBM Cloud上，由200个节点组成，每个节点有8个NvidiaA100 80GB卡，96个VCPU，1.2TB CPU RAM。节点内的GPU卡通过NVLink连接，卡对卡带宽为600GBps，节点通过两条100Gbps以太网链路连接在一起。单根I/O虚拟化(SR-IOV)接口是一种PCIe规范，允许像网卡这样的硬件在PCIe硬件功能之间分离资源访问。基于SR-IOV的TCP/IP栈，提供了120Gbps的可用带宽(对于11B模型，观察到峰值网络带宽利用率为32Gbps)。

“我们想在GPU上投入更多，而不是网络硬件，”IBM研究中心的发明人Raghu Ganti说，他正在研究缩放基础模型。

该GPU系统自5月以来一直在运行，并配置了红帽OpenShift容器平台来运行AI工作负载。该团队正在为大型AI模型的端到端训练、微调和推断构建一个生产准备软件堆栈。

IBM认为，这种方法是业界第一个在标准以太网中使用Kubernetes和PyTorch的FSDPAPI的模型中实现高达110亿个参数的扩展效率的方法。这将允许研究人员和组织以一种更具成本效益和可持续性的方式在任何云中训练大规模模型。2023年，联合团队的目标是继续扩展这项技术，以处理更大的模型。