ATC 2024 | 快手开源大模型长序列训练加速技术,性能大幅超越 SOTA 方案

news/2024/8/27 14:51:04/ 标签: 开源, 语言模型

导读

在深度学习领域,训练大型语言模型(LLMs)一直是一项极具挑战性的任务,它不仅需要巨大的计算资源,同时对内存的消耗也非常巨大。近期,快手大模型团队提出了创新的方法,包括感知流水并行的激活值卸载以及计算-内存均衡的检查点策略,该方法旨在无损地加速大型语言模型的训练过程。目前论文入选ATC 2024大会,代码已开源

图片

  • 论文题目:

    Accelerating the Training of Large Language Models using Efficient Activation Rematerialization and Optimal Hybrid Parallelism

  • 论文地址:

    https://www.usenix.org/conference/atc24/presentation/yuan

  • 代码地址:

    https://github.com/kwai/Megatron-Kwai

论文核心贡献:

1. 感知流水并行的激活值卸载(Pipeline-Parallel-Aware Offloading):针对流水并行产生的激活值,设计卸载和重载的调度策略,充分利用主机内存(主存)来存储激活,且时间开销可以忽略不计。

2. 计算-内存均衡的检查点策略(Compute-Memory Balanced Checkpointing):平衡内存成本和计算成本,以达到帕累托最优,在激活值大小和计算效率之间达到一个实用的平衡点。

3. 性能建模与并行配置调优:针对混合并行配置(张量并行、上下文并行、流水并行、数据并行)搜索空间大的问题,提出了一种性能建模方法,通过测量模型相关元信息、集群相关元信息,就能求解出最优并行配置。

图片

实验结果令人瞩目。以 175B 模型、上下文窗口大小为 32,768 为例,在 256 个 NVIDIA H800 GPU 上,所提出的方法显著地将 MFU(Model FLOPs Utilization)从 32.3% 提高到了 42.7%

一、背景介绍

大模型训练的一个高效实现是 Megatron-LM 框架里的四维混合并行:张量并行(Tensor Parallelism,TP)、上下文并行(Context Parallelism,CP)、流水并行(Pipeline Parallelism,PP)和数据并行(Data Parallelism,DP)。丰富的并行维度带来了不错的性能收益,但也带来了两项挑战。

一、激活值显存瓶颈

随着序列长度增加,每张显卡上的激活值大小等比例增大:例如 175B 模型 32k 序列长度,在不引入跨机 TP 和跨机 CP 的前提下,无论采用何种并行方案,每张显卡上的激活值大小至少是 171.5 GB,这超过了当前任何显卡的显存大小;如果引入跨机 TP 或跨机 CP,由于网络通信速度比 NVLink 低一个数量级,因此通信时间会大幅增加,训练吞吐断崖式下降。

二、并行配置调优困难

混合并行的五个参数(张量并行数 t、上下文并行数 c、流水并行数 p、数据并行数 d、pipeline stage 层数 l)组合空间庞大,例如用 192 卡训练 175B 模型有 287 种能够整除的并行配置。如果手工调优并行参数,则需要消耗大量人力和机器时间,而且当卡数变化、global batch size 变化后需要重新调优。

针对这些问题,论文提出了种减少显卡上激活值的方法,不仅显著提升了长序列训练吞吐,而且提升了常规序列长度的训练吞吐。还提出了一种最优并行配置求解方法,通过测量少量基本性能数据来构建 cost model,使得在任意规模的集群上都能迅速找到最优并行配置。

二、方法概览

感知流水并行的激活值卸载

流水并行分为三个阶段:warm-up 阶段只包含前向计算,steady 阶段前向计算和反向计算交替进行,cooldown 阶段只包含反向计算。由于流水并行的特点,前向计算产生的激活值不会立即被反向计算使用,而这些激活值会持续占用显存。在这段间隔内将激活值卸载到主存,能够减轻显存压力。

感知流水并行的调度方案遵循两个原则:

1) offload 在每个 micro-batch 前向计算结束之后尽快开始;

2) reload 在前一个 micro-batch 反向计算开始时启动。

该方案以 pipeline stage 为调度粒度,不受个别算子计算速度差异、个别激活值大小差异的影响,因此计算与传输能够充分重叠,避免计算与传输相互等待。

整体调度方案如图所示:

图片

在具体实现上,作者采用了以下技术来提高性能:

1)双工传输:在 steady 阶段,卸载和重载同时进行,这有助于利用 PCIe 双向带宽缩短传输时间,减小传输时间无法被掩盖的可能性。

2)乒乓重载:在 GPU 上开辟两块缓冲区,一块缓冲区作为 reload 的目标地址,另一块缓冲区以零拷贝的方式构建激活值提供给反向计算使用;在下一次调度时,两块缓冲区的角色会交换。

3)传输增强:绑定 NUMA(Non-Uniform Memory Access)节点,并使用锁页内存(page-locked memory)来提升传输性能。

主机与设备之间的传输会占用显存带宽,从而影响计算速度,因此传输的激活值越少越好。通过对显存大小建模,计算出最小卸载比例 α,将其余 1 - α 留在显存上,从而最小化传输量和主存占用。

计算-内存均衡的检查点策略

为了减小激活值大小,传统方法是 full checkpointing,即重算每层 transformer layer。这个检查点策略将显存大小减到了最低(仅保存输入张量),而计算代价是增加一倍前向计算量,即 1/3 整体计算量。
本文提出了一种重算代价与收益更加平衡的方法。通过枚举每个激活值所需重算的层,并进行实际测量,得到每个激活值的重算代价。然后枚举所有可能的重算方案,将激活值大小和计算代价画在同一张图上,得到所有重算方案的帕累托包络线。Compute-Memory Balanced Checkpointing 选取了包络线的拐点,将激活值大小的系数从 37.3 节省到 22.7(节省 39%),重算代价仅占总体计算的 1.5%。

具体重算的层如下图所示:

图片

性能建模与并行配置调优

Few-shot 性能建模是一种用比较低的测量代价获取比较准确的时间估计的方案。它将每轮训练迭代的时间拆解成基本性能数据,并且考虑了流水线气泡、算子实现差异、重算时间、通信时间、通信对计算的影响等因素。

图片

基本性能可分为与模型相关性能、集群固有性能两类。模型相关性能包括前向计算、反向计算、重算的时间等,需要为每个模型分别测量;集群固有性能包括 memcpy 传输带宽、DP 传输带宽等,测量一次后可用于所有模型。所有性能指标的测量次数和测量时间如下:

图片

有了这些基本性能数据,预测训练迭代时间无需额外测量。为了得到最优训练配置,枚举所有并行配置,选取预测时间最短的配置即可,求解时间小于 0.001 秒。

三、实验设置

硬件:

实验使用了 32 个节点,每个节点 8 张 NVIDIA H800,主存 1 TB。节点内用 NVLink 连接,节点间用 100Gbps 网卡连接。

软件:

参与对比的最新 Megatron-LM 是 2024.01.01 版本(译者注:ATC '24 截稿日期是 2024.01.16)。Baseline 代码在 Megatron-LM 的基础上做了改进,包括更快的上下文并行实现、更快的 RoPE 实现等。论文代码则在 Baseline 上增加了新提出的 offloading 和 checkpointing 技术。

模型:

参与测试的模型包括 Llama-65B、Llama2-70B、Llama-175B,上下文长度为 4k ~ 128k,默认 global batch size 为 256。其中 Llama2-70B 是 GQA(grouped query attention)模型,其余两个模型是 MHA(multi-head attention)模型;Llama-175B 是指在 Llama-65B 基础上扩大层数和维度的模型。

四、实验结果

性能建模的准确性

采用控制变量来衡量各种并行参数对性能模型的影响。图 8(a)(b)(c) 显示性能模型对于不同的 t 和 c 是准确的。图 8(a)(e)(f) 表明它对不同的 p 和 l 具有鲁棒性。图 8(a)(g) 表明性能模型对所有三种检查点方法都保持正确性。图 8(a)(d) 说明该模型可以适应不同的全局批量大小。在所有这些情况下,实测时间 T 与建模时间 T_model 之间的差异不超过 2.0%。

图片

端到端性能对比

得益于性能建模的准确性,端到端性能对比节省了调优并行配置的时间,且更令人相信 baseline 已调整到最优性能。

实验结果表明,一是 baseline 比最新 Megatron-LM 具有更高的性能和更强的鲁棒性。Baseline 使用了更节省显存和更快速的代码实现,这也是能够鲁棒地进行长序列训练的基础。

二是在卸载和平衡检查点的帮助下,有更大的空间来权衡各种并行配置。与 baseline 相比,性能得到了显著提升。例如,在 256 个 NVIDIA H800 GPU 上,对于上下文窗口大小为 32,768 的 Llama-175B 模型,该方法将 MFU(Model FLOPs Utilization)从 32.3% 提高到 42.7%。

图片

集群规模扩展能力

当集群规模的变化时,手工调优的并行配置很难适配这种变化。常见的手工调优方法是针对特定的卡数进行调优,当卡数变化时等比例地增加或减少数据并行数,这种方法简称 DP 扩展法(data parallel scaling)。这可能导致 global batch size 不满足运行要求;即便 global batch size 满足要求,该并行配置也不一定是新集群规模下的最优配置。

性能模型能够完全适应集群规模变化。例如训练 Llama-65B 模型,上下文长度 4096,global batch size 范围限制在 256 ± 16:给定 24 个节点时,按照性能模型求解出来的最优配置运行,训练吞吐达到了 1.80e5 TPS(Tokens Per Second),而 DP 扩展法只能利用上 20 个节点,吞吐预期只有 1.48e5 TPS。即便限制在 20 个节点,也应该用性能模型求解的最优配置,它达到了 1.59e5 TPS,优于 DP 扩展法。

图片

正确性

从头训练 Llama2-70B 模型,使用 Pile 数据集,上下文长度 4096,开启所有 4D 混合并行技术(TP、CP、PP、DP),以及 offloading 和 checkpointing。实验结果表明,论文提出的训练系统与最新 Megatron-LM 的 loss 曲线一致,说明训练系统与 GQA 和所有 4D 混合并行技术兼容,且没有损害模型的性能。

图片

五、总结

本文提出了两种激活值重建方法,包括感知流水并行的激活值卸载(Pipeline-Parallel-Aware Offloadin),该方法最大程度利用了主存来存储激活值,以及计算-内存均衡的检查点策略(Compute-Memory Balanced Checkpointing),该策略寻求激活值大小和计算效率之间的实际平衡。针对庞大的并行参数搜索空间,提出了一种最优并行配置求解方法,通过测量少量基本性能数据来构建性能模型,从而全面搜索最优的参数组合。

论文代码附带有运行脚本和 Docker 镜像,以确保实验结果的可复现性。

代码在 GitHub 上公开可访问,以促进进一步的研究。


http://www.ppmy.cn/news/1475889.html

相关文章

数模打怪(五)之相关系数

一、什么是相关系数 相关系数:用来衡量两个变量之间的相关性的大小。 根据数据满足的不同条件,选择不同的相关系数进行计算和分析。 两种最为常用的相关系数:person相关系数和spearman等相关系数。 二、Person相关系数 1、什么是Person相…

哈希表(知识点+leetcode)以及字符串哈希

文章目录 一、什么是哈希表二、哈希表常见结构介绍leetcode经典例题242 有效的字母异位词思路编程 349 两个数组的交集思路编程 1 两数之和思路编程 454 四数相加II思路编程 字符串哈希前言思路编程 一、什么是哈希表 哈希表是散列表,就是通过关键码值而直接进行访…

16_Shell好用工具:sed

16_Shell好用工具:sed 零、语法解析 sed [选项参数] [模式匹配/sed命令] 文件 命令说明aadd,新增iinsert,新增cchange,修改ssubstitute,替换ddelete,删除pprint, 打印 通常与 -n 连用 一、增(…

五、 计算机网络(考点篇)

1 网络概述和模型 计算机网络是计算机技术与通信技术相结合的产物,它实现了远程通信、远程信息处理和资源共享。计算机网络的功能:数据通信、资源共享、管理集中化、实现分布式处理、负载均衡。 网络性能指标:速率、带宽(频带宽度或传送线路…

Lua协程(同步的多线程)

1.coroutine.create( func ) 创建一个协程,返回co(coroutine),参数是一个函数,当调用resume时就唤醒co并调用函数 2.coroutine.resume(co, 函数参数们) 启动协程co并传入协程调用函数的参数,可以带回协程…

PHP恋爱话术微信小程序系统源码

💖恋爱高手的秘密武器!恋爱话术微信小程序,让情话信手拈来✨ 💭【开场白:恋爱路上的甜蜜助手】💭 还在为跟心仪的TA聊天时找不到话题而尴尬?或是担心自己说的每句话都显得那么“直男/女”&…

zookeeper和Kafka消息队列群集部署

消息队列概念 什么是消息队列 消息(Message)是指在应用间传送的数据消息队列(Message Queue)是一种应用间通信方式解决方法,确保消息的可靠传输 消息队列的特征 存储 将消息存储在某种类型的缓冲区中,…

电脑如何快速删除相同的文件?分享5款重复文件删除工具

您有没有发现最近电脑运行速度变慢了?启动时间变得更长,甚至完成简单任务也难以如常?这可能是因为重复文件堆积所致。我们发现,清理或移动这些重复的文件和文件夹可以产生惊人的效果。通过删除不必要的重复文件和垃圾文件&#xf…

传输层协议之UDP

1、端口号 我们在应用层创建的套接字,是需要通过bind()接口绑定我们的IP地址与端口号的,这是因为数据从传输层向上交付到应用层时,需要用端口号来查找特定的服务进程。一般在网络通信时,用IP地址标识一台主机,用端口号…

一文学会鉴别“套壳”ChatGPT模型

一文学会鉴别“套壳”ChatGPT模型 随着ChatGPT等明星模型的诞生,市场上也开始出现一些“套壳”现象,即部分模型表面标榜原创或先进,实则在核心算法上与知名模型高度相似。作为技术探索者,如何拨开迷雾,识别这些“李鬼…

蓝桥杯14小白月赛题解

直接输出pi/ti,for遍历 #include <iostream> using namespace std; #define int long long int a,b,c ; double t1.00; signed main() {cin>>a;int an0;for(int i1;i<a;i){cin>>b>>c;if(t>c*1.00/b){tc*1.00/b;ani;} }cout<<an<<e…

MYSQL--第八次作业

MYSQL–第八次作业 一、备份与恢复 环境搭建&#xff1a; CREATE DATABASE booksDB; use booksDB;CREATE TABLE books ( bk_id INT NOT NULL PRIMARY KEY, bk_title VARCHAR(50) NOT NULL, copyright YEAR NOT NULL );CREATE TABLE authors ( auth_id INT NOT NULL PRI…

老物件线上3D回忆展拓宽了艺术作品的展示空间和时间-深圳华锐视点

在数字技术的浪潮下&#xff0c;3D线上画展为艺术家们开启了一个全新的展示与销售平台。这一创新形式不仅拓宽了艺术作品的展示空间&#xff0c;还为广大观众带来了前所未有的观赏体验。 3D线上画展制作以其独特的互动性&#xff0c;让艺术不再是单一的视觉享受。在这里&#x…

大数据之路 读书笔记 Day6 离线数据开发之数据开发平台

回顾 Day5 数据同步遇到的问题与解决方案Day4 数据同步 1. 统一计算平台 1.1 MaxCompute概述 MaxCompute&#xff08;原名 ODPS&#xff0c;Open Data Processing Service&#xff09;是阿里云提供的一种快速、完全托管的EB级数据仓库解决方案。它为用户提供了海量数据存储和实…

STM32智能无人机控制系统教程

目录 引言环境准备智能无人机控制系统基础代码实现&#xff1a;实现智能无人机控制系统 4.1 数据采集模块 4.2 数据处理与控制算法 4.3 通信与网络系统实现 4.4 用户界面与数据可视化应用场景&#xff1a;无人机管理与优化问题解决方案与优化收尾与总结 1. 引言 智能无人机控…

[终端安全]-6 移动终端之应用程序安全

笔者在终端安全专题前面的文章中介绍了移动终端硬件安全和操作系统安全&#xff0c;本文主要介绍移动终端应用安全。在本文最前面&#xff0c;笔者想先解答一位朋友的疑问&#xff0c;为什么需要费心打造一个完整的面面俱到的安全体系&#xff1f; 1 移动终端安全的重要性 移…

C++——类和对象(上)

文章目录 一、类的定义1.类定义格式2.访问限定符3.类域 二、实例化1.实例化概念2.对象⼤⼩ 三、 this指针 一、类的定义 1.类定义格式 与定义结构体类似 class ST {//成员变量int val;//成员函数void print(){cout << val << endl;}};class为定义类的关键字&…

P2p网络性能测度及监测系统模型

P2p网络性能测度及监测系统模型 网络IP性能参数 IP包传输时延时延变化误差率丢失率虚假率吞吐量可用性连接性测度单向延迟测度单向分组丢失测度往返延迟测度 OSI中的位置-> 网络层 用途 面相业务的网络分布式计算网络游戏IP软件电话流媒体分发多媒体通信 业务质量 通过…

【机器学习】Exam4

实现线性不可分logistic逻辑回归 我们目前所学的都是线性回归&#xff0c;例如 y w 1 x 1 w 2 x 2 b y w_1x_1w_2x_2b yw1​x1​w2​x2​b 用肉眼来看数据集的话不难发现&#xff0c;线性回归没有用了&#xff0c;那么根据课程所学&#xff0c;我们是不是可以增加 x 3 x…

【Linux】Vim 使用教程

Linux - Vim Vim 是一款在 Linux 系统中广泛使用的文本编辑器&#xff0c;它是 Vi 编辑器的升级版。Vim 不仅功能强大&#xff0c;而且可高度定制化&#xff0c;是许多程序员和系统管理员的首选工具。以下是 Vim 在 Linux 系统中的安装、配置和使用过程的详细讲解。 附注&…