开发大型语言模型需要投入大量时间和 GPU 资源,这直接导致成本高昂。模型越大,这些挑战就越明显。 最近,Yandex 推出了一种新的解决方案: ,这是一种开源工具,有望通过大幅减少 GPU 资源消耗和训练时间,彻底改变 LLM 训练。在涉及 700 亿个参数的模型的预训练场景中,使用 YaFSDP 可以节省大约 150 个 GPU 的资源。这意味着每月可节省约 50 万美元至 150 万美元,具体取决于虚拟 GPU 提供商或平台。 YaFSDP Yandex 已将 YaFSDP 公开发布于 。 GitHub 分布式法学硕士培训的挑战 在多个 GPU 上训练 LLM 涉及复杂的操作,这会导致效率低下和内存消耗高。主要问题之一是需要在 GPU 之间发送和接收大量数据。例如,在典型的 all_reduce 操作中,必须传递的梯度数据量是网络参数的两倍。对于 Llama 70B 模型,这意味着每次迭代要传输 280 GB 的数据。 此外,权重、梯度和优化器状态在 GPU 之间重复,导致巨大的内存负载。Llama 70B 模型和 Adam 优化器需要超过 1 TB 的内存,远远超过大多数 GPU 的典型 80 GB 内存容量。这种冗余严重减慢了训练过程,并且通常使得即使中等大小的模型也难以装入 GPU 内存。 YaFSDP 介绍 Yandex 的 YaFSDP 为这些挑战提供了高效的解决方案。通过专注于优化内存消耗和消除通信瓶颈,YaFSDP 提高了 LLM 训练的效率。它通过分片层而不是单个参数来工作,保持高效的通信并避免冗余操作。此外,YaFSDP 为所有必需数据预先分配缓冲区,确保 Torch 分配器不会引入低效率。 YaFSDP 通过利用两个缓冲区来存储中间权重和梯度,其中奇数层使用一个缓冲区,偶数层使用另一个缓冲区。 不同层的权重存储在同一个内存中。如果层具有相同的结构,它们将始终相同。确保当您需要第 X 层时,缓冲区包含第 X 层的权重至关重要。所有参数都将存储在缓冲区内的相应内存块中。 内存消耗 在训练期间,主要的内存消耗者是权重、梯度、优化器状态、缓冲区和激活。YaFSDP 通过优化这些元素的存储和访问方式显著减少了内存消耗。 :这些取决于进程数量,随着进程数量的增加,它们的内存消耗趋于零。通过将这些组件分片到 GPU,YaFSDP 可以最大限度地减少重复,从而减少内存使用量。 权重、梯度和优化器状态 消耗恒定量的内存并在计算期间存储中间值。 缓冲区 取决于模型大小和每个 GPU 处理的令牌数量。 激活 激活检查点 激活检查点是一种在正向传递期间仅存储必要的激活并在反向传递期间重新计算它们的技术。这显著减少了内存占用,因为只存储了必要的数据。例如,在训练批大小为 8192 个标记的 Llama 2 70B 模型时,激活存储可以从超过 110 GB 减少到仅 5 GB。 然而,这种方法引入了额外的计算开销,YaFSDP 通过不对某些层使用激活检查点来避免这种开销,这可以通过内存优化来实现。 沟通优化 YaFSDP 通过确保仅在必要时传输数据并将通信与计算重叠来提高 GPU 通信效率。它利用 CUDA 流来有效地管理并发计算和通信。 该工具使用两个流:计算流和通信流。事件同步这些流,确保操作按正确的顺序执行,而不会引入死锁。 第三层的正向传递直到 all_gather 操作完成才会开始(条件 1)。同样,第三层的 all_gather 操作直到使用相同缓冲区的第一层的正向传递完成才会开始(条件 2)。由于此方案中没有循环,因此不可能出现死锁。 实验结果和性能提升 YaFSDP 的实施已显示出训练效率的显著提升。在一个具有 700 亿个参数的模型的预训练场景中,YaFSDP 能够节省大约 150 个 GPU 的资源。这意味着每月可节省大量成本,根据虚拟 GPU 提供商或平台的不同,每月可节省 50 万美元至 150 万美元不等。 与 FSDP 等现有方法相比,YaFSDP 将训练时间缩短了多达 26%,并优化了内存使用情况,从而可以更高效地训练更大的模型。 Yandex 已将 YaFSDP 公开发布于 . ML 工程师可以利用此工具来提高其 LLM 培训流程的效率。通过开源 YaFSDP,Yandex 旨在促进 AI 社区的创新和协作,使开发人员能够更快、更经济高效地训练模型。 GitHub YaFSDP 代表了 LLM 训练的重大进步。解决内存消耗和通信效率低下的关键挑战可以更快、更高效地训练大型语言模型。