2023年百度大模型研究笔记(二) 大规模计算集群优化 全球速讯

2023-05-31 10:38:33

来源:中信建投

1、模型结构

1.1、输入长度(Context Window)

Context Window 延长拓展大模型的能力圈。5 月 11 日,Anthropic 宣布将 Claude AI 的输入长度(Context Window)从 9K 扩展至 100K tokens1,对应约 75000 个英文单词。根据《How many words do we read per minute? A review and meta-analysis of reading rate》,人类阅读 100K tokens 的文本内容大约需要 5 个小时以上,而 Claude 能显著缩短这一时长,并且在分辨长文本差异的速度上,Claude 也具备优势,例如 Anthropic 团队将《了不起的 盖茨比》(72K tokens)中修改了一处句子,Claude 仅用时 22 秒定位文本的变化。通过扩展 Context Window, Claude 具备处理长文本、长语言等能力,提升其在协同办公等领域的能力。

引入局部注意力机制的研究最早见于 ALiBi。2021 年 8 月,华盛顿大学、Meta、Allen 人工智能研究所的团 队提出了 ALiBi(Attention with Linear Biases)编码,延续相对位置编码思路,但 position embedding 并没有加在 word embedding 上,而是加上 query 距离成正比的惩罚值,最终实现更稳健的输入长度线性外推表现。


(资料图片仅供参考)

KERPLE 在 ALiBi 基础上改进,取得效果提升。2022 年 5 月,CMU 和普林斯顿大学的研究团队将 Attention 计算公式进一步调整,并引入1 , 2两个可学习参数,从 ALiBi 的 − | − |调整为 { + −1 | − | 2, 1 > 0,0 < 2 ≤ 2 + − 1 log(1 +2 | − |) , 1 , 2 > 0 最终 KERPLE 实现了对 ALiBi 效果的改进,并在 OpenWebText2、GitHub、ArXiv 数据集上的长文本输入 方面领先此前的编码方式。

微软研究团队于 2022 年 12 月提出 XPOS(衍生自 RoPE),实际上对 Q、K 做了非对称变化,导致最终 结果引入超参数的变化,且也是以来相对距离 − ,而非此前 ALiBi 的| − |,由于是单向模型,因此只需 要考虑 ≥ 的情况,并设置超参数在 0~1 的范围变化即可(相当于指数衰减)。

局部注意力机制是对文本的人为截断,且在非语言任务的泛化能力上可能存在问题。因此,一些研究寻求 其他路径。回归输入长度延长的思考,本质上这是一个训练长度和预测长度不一致的问题,其可能是 1)预测时 用到没训练过的位置编码;2)预测时注意力机制所处理的 token 数量远超训练时的数量。 针对位置编码未被训练,《Randomized Positional Encodings Boost Length Generalization of Transformers》提 出在训练时随机采样位置序列或者均匀取点,实现每个位置编码都被训练到。最终实验表明随机位置训练 有效地提升了长度外推的表现。

针对注意力机制处理 token 数量,一个思路是引入缩放注意力,使得注意力机制在处理大量 token 时,注 意力的熵随着长度的变化保持平稳。其中是训练长度,是预测长度。 (,, ) = ( √ ) (,, ) = ( log √ ) 总结来看,局部注意力机制和缩放注意力/随机 Token 处理策略都是某种程度上基于局域信息,也就是如果 下游任务不具有局域性,输入长度的延长属性是否还成立?Google 团队于 2022 年 7 月3提出,通过引入长度泛 化性基准测试(Chomsky Hierarchy Evaluation Benchmark),简单理解即测试时输入序列的长度可以是无限的,实验表明 Transformer 的在多个下游任务的表现上线性延展性能弱于 RNN、LSTM 模型。

在随机位置 训练之下,不同位置编码设计下的 Transformer 都有明显提升,这反映位置编码设计与线性延展性关联度不高。

除了对局域性的处理,另一些思路则沿着降低计算复杂度,从而实现给定内存/算力资源下文本处理能力 的扩展,例如 Google 和 UCB 团队于 2020 年 1 月提出的 Reformer4,以及 Google 研究团队于 2020 年 3 月提出 的 Routing Transformers5。Reformer 主要是 1)引入可逆层,降低模型内存;2)FFN 层分块并行处理,降低内 存消耗;3)局部敏感哈希替代 dot-product attention,降低计算和内存复杂度,最终将计算复杂度从( 2 )降至 ( log )。Routing Transformers 则结合内容稀疏注意力机制和局部注意力机制,将计算复杂度从( 2)降至 ( 1.5)。

FlahAttention 优化内存 I/O 瓶颈,大幅提升 token 处理能力。但斯坦福大学和纽约州立大学团队于 2022 年 5 月研究指出,此前的研究尽管从结果上优化了内存/计算复杂度并降低了时长,但往往聚焦于特定环节,对整 体的优化幅度有限。这主要是由于大规模计算集群的资源瓶颈主要是内存读取速度,因而 FlashAttention 主要针对内存控制优化。FlashAttention 主要通过 1)调整注意力计算机制,将输入分成块,并在输入块上进行多次传 递,从而逐步执行 softmax reduction;

2)存储前向传递的 softmax 归一化因子,在后向传播中快速重新计算片 上注意力,这比从 HBM 中读取中间注意力矩阵的标准方法更快。该研究在 CUDA 中实现 FlashAttention,以 达到对内存访问的细粒度控制,并将所有注意力操作融合到一个 GPU 内核中。即使由于重新计算导致 FLOPs 增加,但其运行速度更快(在 GPT-2 上比传统算法速度提升 7.6 倍)并且使用更少的内存(序列长度线性), 主要是因为降低了 HBM 访问量。

在此之上,研究团队提出块稀疏注意力(Block-Sparse FlashAttention)算法,实现对 FlashAttention 的进一 步加速。且块稀疏注意力算法下 transformer 序列处理准确度在 64K 下达到 63.1%。斯坦福大学和蒙特利尔大学的研究团队于 2023 年 2 月进一步提出 Hyena 架构,涵盖 1)Element-wise 乘法 门控;2)长卷积,最终实现线性延展性性能损失的最小化,同时大幅节约计算成本。Hyena 在百万级 token 长 度下实现超越 50%的准确率,且相比 transformerFLOPs 下降 20%,在 8K 文本处理速度下较 FlashAttention 提升 2x,在 64K 文本下处理速度较 FlashAttention 提升 100x,标准方法下模型无法处理该内容长度。

2、预训练

2.1、并行化

随着参数规模扩大,显存压力导致模型训练难度提升。学界主要有两种思路:1)数据并行化;2)模型并 行化。数据并行化是将数据集并行分布在不同设备上进行训练,从而实现 batch size 规模的线性化提升,但难 点在于对显存使用效率不高,因此一般适用于中小规模的模型训练。模型并行化则是将不同层切分开来,由不 同卡负责计算模型的一部分,这导致额外的通信操作,主要特点是显存使用效率高,但通信效率低。

Google 和丰田研究院的团队在 2019 年 9 月提出 ALBERT7,即通过共享参数防止参数量随着网络深度的加 深而增加,同时分解 embedding 层参数使得增加隐层大小时而不用改变词嵌入的大小。这在一定程度上解决了 数据并行下模型规模难以扩张的问题,但由于 ALBERT 基于 Encoder 编码器、Gelu 函数,因此其难以泛化至 GPT 等 Decoder 模型架构中。

微软研究团队于 2019 年 10 月提出 ZeRO,通过分片(partition)显著优化显存和通信花费。值得一提的 是,后续 DeepSpeed 超大规模训练工具正是基于 ZeRO 为代表的一系列工作。在《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》中,研究团队探讨了内存显存的结构,并将将模型训练阶段每张卡中 显存内容分为两类:1)模型状态(model states);2)剩余状态(residual states),而模型状态占用了主要显存, 因为深度学习中应用比较广泛的 Adam 优化器涵盖了参数梯度、梯度的一阶动量和二阶动量,并且在混合精度 训练下需要存储 FP16 的模型参数、梯度,FP32 的 Adam 状态(模型参数、梯度备份、梯度的一阶动量和二阶 动量)。

定量看,如果单模型参数量为,则实际需要 16的存储空间(其中 75%来自 Adam 优化器)。因此研 究团队提出通过分片(partitions)优化 Adam 带来的模型状态显存占用,通过动态通信策略提升通信效率,+ 下模型通信量和标准的数据并行一致。

微软研究团队后续于 2021 年 1 月、2021 年 4 月发布 ZeRO-Offload 与 ZeRO-Infinity。其中 ZeRO-Offload 就 是引入相对 GPU 显存更廉价的 CPU 内存,但尽可能避免 CPU 通信对整体系统的拖累;ZeRO-Infinity 相较于 Offload 聚焦单卡场景,更适用于超大规模训练场景(业界应用),资源利用率达到 40%水平。Meta 研究团队于 2021 年 7 月提出 FSDP8,对标 ZeRO。FSDP 在此前 DDP9(分布式数据并行)基础上做了 优化,结合参数切分(parameter sharding),实现训练效率提升。

ALBERT、ZeRO、FSDP 主要是从数据并行的角度,结合模型并行的思路进行改进,而另一些研究从模型 并行出发,例如英伟达的 Megatron、Google 的 Gpipe、Kakao 的 torchpipe 等。其中,英伟达团队于 2019 年 9 月 提出 Megetron-LM,基于模型并行基础对 MLP 和 Attention 做切分,最终实现单机 8 卡达到 77%的资源利用率 (等效于 6 卡),模型和数据混合并行下 512 块 GPU 资源利用率达到 74%。英伟达团队于 2021 年 4 月在此基 础上提出结合数据并行方法,实现资源利用率的提升,在 3072 块 GPU 上训练万亿参数模型,计算量达 502 petaFLOP/s,利用率达 52%,较 ZeRO-3 实现更强的效率提升。

Google 研究团队于 2018 年11月提出 Gpipe,通过流水线并行扩展深度学习的训练负载。GPipe 引入流水 并行策略来缓解顺序执行的性能问题,把 mini-batch 细分为多个更小的 micro-batch,同时通过 recompute 降低内 存。后续 Gshard、GSPMD、GaLM 引入自动分片、稀疏激活等技术,实现性能提升。微软和斯坦福大学、CMU 的研究团队于 2018 年6月提出PipeDream,通过结合模型、数据、流水并行解 决数据并行带来的大量通信成本。

2.2、大规模计算集群优化

国内方面,百度与英伟达联合开发 AI 计算集群,采用百度自研 X-MAN 架构 4.0 及英伟达 A100-80GB、 Infiniband 网卡,百度硅谷 AI 实验室(SVAIL)于 2017 年 2 月将 HPC 领域的 Ring AllReduce 迁移至深度学习 领域,解决多 GPU 环境的通信瓶颈问题,目前 Ring AllReduce 已经成为行业较为普遍应用的算法。腾讯 4 月发 布新一代高性能计算集群(High-Performance Computing Cluster),采用腾讯云星星海自研服务器,搭载英伟达 H800 GPU,服务器之间采用 3.2T 互联带宽。

阿里巴巴研究团队 2020 年 2 月提出了 EFLOPS 集群,主要优化了大规模计算集群的通信网络问题。过往的 大规模计算集群的通信网络瓶颈主要来自 1)网卡;2)PCIe 总线;3)不均匀的 QoS 分布。阿里提出为每块 GPU 配置网卡,并使用基于 CLOS 架构的数据中心,通过对端口进行编号的方式构建一个层次化的拓扑结构。实验 表明,最终多线程带宽最高实现 40Gbps 左右水平。

后续 Cisco 的研究团队于 2021 年分别提出 Ruta11/NetDAM12,前者通过引入 3D-Torus 拓扑结构,移除交换 机,实现近乎为零的丢包和低于 200ms 的延迟,较此前的研究有明显提升。NetDAM 主要解决网络拥塞控制 问题,研究团队指出丢包本质是一种传输的损失,而拥塞控制算法本质上是为了降低损失率,而丢包损失由不 确定性影响,所谓的拥塞控制算法就是估计损失的数学期望,并尽可能优化方差和均值。因此,后续的优化方 向主要是确定性网络(降低损失率),准确测量(优化估计从而便于控制)。

在确定性网络方向,阿里巴巴于 2022 年 SIGCOMM 发布 µFAB 和 Solar;在准确测量方向,阿里巴巴于 2019 年 SIGCOMM 发布 HPCC,提出通过 in-network telemetry 实时监测 buffer,缓解 TCP incast 问题。后续哈 佛、USC 等研究团队于 2020 年 7 月提出《PINT: Probabilistic In-band Network Telemetry》,对 HPCC 进行改进, 实现效果提升。而 Google 在 2020 年发布《Swift: Delay is Simple and Effective for Congestion Control in the Datacenter》,提到 DCTCP、PFC、DCQCN、HPCC 使用来自交换机的显式反馈,以保持低延迟,但由于 HPCC 实时监测带来和交换机的紧密协同,增加部署和维护难度。Google 的思路则是降低对网络硬件的依赖,解耦主 机和网络。

Cisco 于 2021 年 10 月提出 NetDAM,思路与 Google 类似,主机侧解耦计算域与 I/O 域内存的方式, 使得响应延迟具有了确定性,同时通过无锁共享内存实现通信,降低交换网 incast 概率。

2.3、内存置换与Checkpoint优化

内存置换即 Offloading,例如通过 CPU 内存置换 GPU,并通过 CPU 和 GPU 通信实现扩容和成本优化,一 个典型的案例就是此前提到微软团队于 2021 年 1 月提出 ZeRO-Offload,引入相对 GPU 显存更廉价的 CPU 内 存,但尽可能避免 CPU 通信对整体系统的拖累。

另一个训练技巧则是 Checkpoint。Checkpoint 引入的初衷也是为了解决存储空间的瓶颈,因此在 Checkpoint 中存储部分前馈传输的激活点和权重,以及在后馈时计算剩余的激活点和权重。2016 年 4 月华盛顿大学、MIT 研究团队提出《Training Deep Nets with Sublinear Memory Cost》,核心思想是将部分占用内存空间大,但计算消 耗小的 tensor,不进行存储,而在需要时重新计算即可。研究的主要缺陷在于忽略了网络中的异质性。2019 年 11 月法国波尔多大学(Université de Bordeaux)、蒙彼利埃大学(Université de Montpellier)等团队提出针对异 质性网络的优化 Checkpoint 方法。

后续法国国家信息与自动化研究所(Inria Bordeaux)在《Efficient Combination of Rematerialization and Offloading for Training DNNs》将 offloading 与 checkpoint 技巧结合,其中 offloading 是用 内存置换通信需求(避免通信瓶颈),checkpoint 是用内存置换计算时间(降低计算资源消耗)。将两种技术合 并将进一步优化整体性能。

3、下游调试

3.1、参数高效调试(Parameter-Efficient Tuning)

由于 Scaling Law 及 CoT 带来的规模竞赛,当前主流预训练大模型的参数规模普遍在数十亿乃至千亿级别, 这导致对所有参数做精调难度较大,因此学术界提出只微调部分参数的思路,但过往的研究一般存在性能损失 等问题,即微调后模型性能无法与全参数调试的性能相比。微软和 CMU 的研究团队于 2021 年 6 月提出 LoRA13, LoRA 的核心思想是过参数模型存在低内在维度,因此可以通过秩14分解矩阵来间接训练神经网络中的一些密集 层,同时冻结预训练模型权重,降低了存储占用,同时提升训练速度(减少计算量)。

另一种高效训练的模式则是基于成熟模型的回答构建的。斯坦福大学 2023 年 3 月提出 Stanford Alpaca15, 基于 LLaMA 7B 模型,以及 5.2 万条基于人类指令反馈内容(数据来自 ChatGPT)。实验表明,Alpaca 的能力 接近 OpenAI 的 text-davinci-003,但参数规模和训练成本远低于 text-davinci-003。另一方面,Alpaca 也相应继承 了原有模型存在的不足,例如欺骗性、错误价值观等,仍需要进一步的调试16和 alignment(对齐)。需要注意的 是,由于基于 LLaMA 开源模型,其需要遵守相应地非商用协议,即基于 LLaMA 的模型不得用于商业用途。另 外,其调试数据来自 ChatGPT(text-davinci-003),因此基于这些数据训练、调试的模型不得与 ChatGPT 进行 竞争。违反以上协议可能面临相应的法律合规风险。

受 Alpaca 启发,2023 年 3 月 UCB、CMU、Stanford、UCSD、MBZUAI 研究团队推出 Vicuna-13B,即基 于 LLaMA-13B 模型,并基于 7 万条用户分享的 ChatGPT 相关反馈做调试,并实现接近 ChatGPT 90%的性能 (由 GPT-4 评判),接近 Bard,领先 Alpaca、LLaMA 等模型。对 LoRA 的评估方面,SUTD、SMU、西南交通大学、达摩院等研究团队于 2023 年 4 月指出,在复杂任务 上,经过 LoRA 等方法调试后的微缩版模型性能弱于基础模型,但在简单任务上二者表现比较接近,且 LoRA 总体上优于其他调试方法。

3.2、基于人类反馈的强化学习(RLHF)

2023 年 5 月,Stanford 和多伦多大学的研究团队提出 AlpacaFarm,旨在低成本、快速实现 RLHF。AlpacaFarm 主要由三个环节组成:1)基于人类反馈数据的训练;2)相比基准模型的评估;3)对比其他 RLHF 方法。相比 于传统的 RLHF,alpaca 引入基于 LLMs 的人类反馈数据,即基于成熟模型的数据作为“标准答案”,这种设计 下成本能够压缩为传统方法的 1/45。RLHF 主要有四个主要问题,1)如何选择人类反馈的形式,存在打分、排序、文本描述等类型;2)目标 主要包括有用性与无害性;3)使用机制分为训练(基于反馈的模仿学习、联合反馈建模、强化学习),解码 (再排序、反馈调节);4)如何将上述反馈机制模型化。

关于人类反馈机制的形式,排序是一个应用较为普遍的机制,因为打分涉及较多主观因素,标注者的价值 观不同可能导致分数需要归一化处理,而排序则避免了这一问题17。基于同一 prompt 得到不同模型的反馈排序 后,进一步通过 Elo 系统打分归一化。其他反馈形式转化为数值的难度较高,或存在主观性较强等问题,因此 目前主流研究采用排序居多。

关于使用机制方面,John Schulman 在《Reinforcement Learning from Human Feedback: Progress and Challenges》 18提到,SFT 与其等价的 BC(behavior cloning)存在固有缺陷,即训练越充分越容易出现欺骗(即模型并不明 确自己知识的边界),RLHF 则是让模型明确有些问题自己不知道。原理上是因为 SL 训练时只有正反馈,而且 对偏离样本的惩罚较重,RL 多样性更好,因而在面对不知道的问题时,SL 训练充分的模型倾向于回答(胡编 乱造),而非反馈不知道19。RL 方法中,《AlpacaFarm: A Simulation Framework for Methods that Learn from Human Feedback》测试下 PPO 效果最优。

关于建模方面,如 AlpacaFarm 等研究采用基于人类反馈的建模,即构建一个函数使得函数在人类反馈数据 集上与真实人类的反馈程度距离不要相差太远,并通过反馈建模规模化的优化预训练模型,解决了传统方法下 人类反馈带来的高成本问题。但 2022 年 10 月 OpenAI 研究团队指出20,在偏好模型中,使用少于 1000 次比较 仅能导致微小改进,结果近似于随机。此外,采用静态反馈(人类反馈结果不更新)会导改变模型输出分布, 影响模型的一致性和准确性。

因而,学界/业界开始引入基于 AI 反馈的强化学习(RLAIF),例如 2022 年 12 月 Anthropic 研究团队在 《Constitutional AI: Harmlessness from AI Feedback》提出,人类反馈本质上也是基于一系列的原则,因此在监督 学习阶段,研究团队通过 promt 引导模型生成有害性的回答,并将原则以 prompt 形式输入模型,要求模型判断 此前的回答是否违背原则。迭代训练后,将模型面临有害性 prompt 下基于原则生成的无害性回答收集起来,并 用于预训练模型的 finetune。在强化学习阶段,研究团队让此前 SL 训练的模型生成一系列有害性的 prompt 并让 模型在一系列 prompt 中进行有害性排序(类似于 RLHF 的排序),并将 AI 回答构成的有害性排序数据集与人 类生成的有用性数据集结合,共同训练偏好模型,并用于大模型的调试。

Anthorpic 与 OpenAI 等团队主要采用Self AI Feedback,即生成回答与反馈是同一个模型。另一种技术路 径则是引入外部 AI 模型反馈,其优势在于 Self AI Feedback 意味着反馈模型参数规模较大,训练/调试的成本 和时间较多,而外部反馈模型在效率方面具备一定优势。

(本文仅供参考,不代表我们的任何投资建议。如需使用相关信息,请参阅报告原文。)

关键词: