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

news2024/12/23 18:00:43

导读

在深度学习领域,训练大型语言模型(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.coloradmin.cn/o/1926079.html

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈,一经查实,立即删除!

相关文章

【学习笔记】无人机(UAV)在3GPP系统中的增强支持(二)-支持高分辨率视频直播应用

引言 本文是3GPP TR 22.829 V17.1.0技术报告,专注于无人机(UAV)在3GPP系统中的增强支持。文章提出了多个无人机应用场景,分析了相应的能力要求,并建议了新的服务级别要求和关键性能指标(KPIs)。…

TS真的比JS更好吗?

前言 在讨论TypeScript(TS)是否比JavaScript(JS)更好时,我们需要明确“更好”这一概念的上下文和衡量标准。TypeScript和JavaScript在多个方面有着明显的区别,但它们并不是简单的“好”与“不好”的关系&a…

springboot上传图片

前端的name的值必须要和后端的MultipartFile 形参名一致 存储本地

一文搞定:Syncthing多平台文件同步工具安装全攻略

Syncthing是一款开源的文件同步工具,可以通过本地网络或互联网实现多台设备之间的文件同步。与其他同步工具不同,Syncthing强调隐私和安全,确保用户的数据始终处于用户的控制之下。 功能与特点 开源软件: Syncthing是完全开源的&…

base SAS programming学习笔记11(functions)

1.SAS function 分类&#xff1a; 计算描述统计量的函数&#xff1a; 举例如下&#xff1a;avgscoremean(exam1,exam2,exam3) 2.function 基本格式 function-name(argument1,argument2,......<argumentn>&#xff09; argument可以如下&#xff1a;变量名&#xff1b;常…

新手小白的pytorch学习第三弹-------tensor的基本操作

reshape, view, stacking, squeeze(), unsqueeze(),permute()torch.tensor 和 numpy 的 array切片&#xff0c;张量里面获取元素值随机种子 1 导入torch import torch2 reshape() tensor_A torch.arange(1, 11) tensor_Atensor_A.reshape(2, 5) tensor_A.reshape(2, 5)tenso…

C语言 ——— 实用调试技巧(Visual Studio)

目录 Debug 和 Release 的区别 F10 --- 逐过程调试 & F11 --- 逐语句调试 F9 --- 新建/切换断点 & F5 --- 开始调试 shift F5 & ctrl F5 Debug 和 Release 的区别 Debug&#xff1a;通常为调试版本&#xff0c;它包含调试信息&#xff0c;并且不作任何优化…

Unity ColorSpace 之 【颜色空间】相关说明,以及【Linear】颜色校正 【Gamma】的简单整理

Unity ColorSpace 之 【颜色空间】相关说明&#xff0c;以及【Linear】颜色校正 【Gamma】的简单整理 目录 Unity ColorSpace 之 【颜色空间】相关说明&#xff0c;以及【Linear】颜色校正 【Gamma】的简单整理 一、简单介绍 二、在Unity中设置颜色空间 三、Unity中的Gamma…

Vortex GPGPU的硬件代码分析(Cache篇2)

文章目录 前言一、VX_cache.sv代码部分解读2——buffering/initialize1.1 core response buffering与VX_elastic_buffer模块解读1.1.1 VX_pipe_buffer模块解读1.1.1.1 一种握手信号的解释1.1.1.2 世界线收束——VX_pipe_buffer的核心代码解释1.1.1.3 VX_pipe_register模块解读与…

算法015:串联所有单词的子串

串联所有单词的子串. - 备战技术面试&#xff1f;力扣提供海量技术面试资源&#xff0c;帮助你高效提升编程技能,轻松拿下世界 IT 名企 Dream Offer。https://leetcode.cn/problems/substring-with-concatenation-of-all-words/ 如果是第一次接触这个题目&#xff0c;接触滑动…

埋点系统如何统计用户的平均停留时长?

Hello&#xff0c;大家好&#xff0c;欢迎使用Webfunny前端监控和埋点系统。 今天&#xff0c;我们将介绍webfunny的埋点系统如何统计用户的平均停留时长 一、页面beforeLeave事件 当你页面离开的时候&#xff0c;会触发一个心跳检测&#xff0c;但是这个可能不是100%触发&am…

跳表的简单学习

跳表&#xff08;SkipList&#xff09;学习 1. 什么是跳表&#xff1f; 基于“空间换时间”思想&#xff0c;通过给链表建立索引&#xff0c;使得链表能够实现二分查找。 跳表是可以实现二分查找的有序链表。 2. 从单链表到跳表 对于一般的单链表&#xff0c;在其中进行查…

EasyCVR视频技术:城市电力抢险的“千里眼”,助力抢险可视化

随着城市化进程的加速和电力需求的不断增长&#xff0c;电力系统的稳定运行对于城市的正常运转至关重要。然而&#xff0c;自然灾害、设备故障等因素常常导致电力中断&#xff0c;给城市居民的生活和企业的生产带来严重影响。在这种情况下&#xff0c;快速、高效的电力抢险工作…

【PVE】新增2.5G网卡作为主网卡暨iperf测速流程

【PVE】新增2.5G网卡作为主网卡暨iperf测速流程 新增网卡 新增网卡的首先当然需要关闭PVE母机&#xff0c;把新网卡插上&#xff0c;我用淘宝遥现金搞了个红包&#xff0c;花了26元买了块SSU的2.5G网卡。说实话这个价位连散热片都没有&#xff0c;确实挺丐的。稍后测下速度看…

从零开始做题:segmentFlow

题目 解题 import string import binascii def Crack_4B(crc_list):print(-------------Start Crack CRC 4B-------------)#crc_list [0xc0a3a573, 0x3cb6ab1c, 0x85bb0ad4, 0xf4fde00b]#文件的CRC32值列表&#xff0c;注意顺序comment chars string.printableflag 0 for …

文心一言的流式接口数据进行处理 增加属性

需求&#xff1a;需要对文心一言的流式接口数据进行处理 增加属性 return ResponseEntity.ok().header("Access-Control-Allow-Origin", "*").contentType(org.springframework.http.MediaType.TEXT_EVENT_STREAM).cacheControl(org.springframework.http…

Leetcode3201. 找出有效子序列的最大长度 I

Every day a Leetcode 题目来源&#xff1a;3201. 找出有效子序列的最大长度 I 解法1&#xff1a;枚举 全奇数、全偶数、奇偶交替三种情况的最大值即为所求。 代码&#xff1a; /** lc appleetcode.cn id3201 langcpp** [3201] 找出有效子序列的最大长度 I*/// lc codesta…

imx6ull/linux应用编程学习(16)emqx ,mqtt创建连接mqtt.fx

在很多项目中都需要自己的私人服务器&#xff0c;以保证数据的隐私性&#xff0c;这里我用的是emqx。 1.进入emqx官网 EMQX&#xff1a;用于物联网、车联网和工业物联网的企业级 MQTT 平台 点击试用cloud 申请成功后可得&#xff1a;&#xff08;右边的忽略&#xff09; 进入…

【matlab】大数据基础与应用实例

目录 引言 线性回归模型 基本形式 最小二乘法 多元线性回归 线性回归的假设 模型评估 应用 独热编码 原理 应用场景 优点 缺点 数据收集 数据可视化 数据处理与分析 完整代码 引言 线性回归模型 线性回归模型是一种用于预测连续值输出&#xff08;或称为因变…

「iOS」暑假第一周 —— ZARA的仿写

暑假第一周 ZARA的仿写 文章目录 暑假第一周 ZARA的仿写写在前面viewDidLoad 之中的优先级添加自定义字体下载想要的字体添加至info之中找到字体名字并应用 添加应用图标和启动页面 写在前面 暑假第一周留校学习&#xff0c;对于ZARA进行了仿写&#xff0c;在仿写的过程之中&a…