思考一下，联邦学习可以训练大语言模型吗？

机器之心原创

作者：Jiying

编辑：H4O

1、概述

随着大语言模型（Large Language Model，LLM）的火速发展，关于大语言模型对人工智能产业发展的影响引发了越来越多的讨论。一种观点认为，大语言模型的发展摧毁了人工智能初创企业的发展之路，因为大语言模型参数多，所需要的算力规模大，所依赖的训练数据规模也大。大模型、大参数、大数据实际都集中在大的人工智能企业，从而导致初创企业的机会越来越少。另外一种观点则相反，他们认为，大语言模型的发展一定程度成促进了人工智能在多个领域中的广泛发展，例如可以直接在大语言模型的基础上利用私有数据搭建一些垂直领域的大语言模型，可以直接将大语言模型应用在不同的业务场景中等等。

我们在这篇文章中不讨论上述两种观点，我们关注的是第二种观点中提到的应用方式：如何在保证数据隐私性的前提下，利用私有数据训练大语言模型，从而满足在垂直领域中的应用需求？这一点也将是 LLM 促进人工智能初创企业发展的关键。具体的，我们关注联邦学习是否可以用来训练大语言模型？

1.1 大语言模型结构及需要的资源

由于 Transformer 架构具有出色的可并行性和容量，它已经成为开发各种 LLM 事实上的主干网络，使得语言模型的规模达到数千或数亿的参数成为可能。一般来说，现有 LLM 的主流架构可以大致分为三大类型，即编码器 - 解码器（encoder-decoder）、因果解码器（causal decoder）和前缀解码器（prefix decoder），如图 1 所示 [8]。

图 1. 三种主流架构中的注意力模式比较。其中，蓝色、绿色、黄色和灰色的圆角矩形分别表示前缀标记之间的注意力、前缀和目标标记之间的注意力、目标标记之间的注意力和掩码的注意力。

• 编码器 - 解码器架构。Vanilla Transformer 模型是建立在编码器 - 解码器结构上的，它由两个堆叠的 Transformer 块组成，分别作为编码器和解码器。编码器采用堆叠的多头自注意力层对输入序列进行编码，以生成其潜在的表征，而解码器则对这些表征进行交叉注意力，并自动生成目标序列。编码器 - 解码器（如 T5 和 BART）已经在各种 NLP 任务上显示出有效性。到目前为止，只有少数的 LLM 是基于编码器 - 解码器架构建立的，如 Flan-T5。
• 因果解码器架构。因果解码器的结构包含了单向的注意力掩码，以保证每个输入标记只能注意到过去的标记和自己。输入和输出标记以同样的方式通过解码器进行处理。作为该架构的代表性语言模型，GPT 系列模型是基于因果解码器架构开发的。特别是 GPT-3 成功地证明了这一架构的有效性，也显示了 LLM 惊人的内涵学习能力。有趣的是，GPT-1 和 GPT-2 并没有表现出像 GPT-3 那样优越的能力，看来缩放在提高这个模型架构的模型容量方面发挥了重要作用。到目前为止，因果解码器已被各种现有的 LLM 广泛采用为 LLM 的架构，如 OPT、BLOOM 和 Gopher。
• 前缀解码器架构。前缀解码器架构（又称非因果解码器）修改了因果解码器的掩码机制，以便能够对前缀标记进行双向关注，而只对生成的标记进行单向关注。这样，像编码器 - 解码器结构一样，前缀解码器可以对前缀序列进行双向编码，并逐一自动预测输出标记，在编码和解码期间共享相同的参数。与其从头开始预训练，一个实用的建议是不断训练因果解码器，然后将其转换为前缀解码器，以加速收敛，例如，U-PaLM 是由 PaLM 衍生出来的。现有的基于前缀解码器的代表性 LLM 包括 GLM130B 和 U-PaLM。

但是，在大语言模型发展过程中，越来越多的人提出了疑虑，主要集中在以下几个方面：

• 隐私不应受到损害。最先进的模型只能通过通用的黑盒 API 访问。这些 API 在向集中式 LLM 供应商传输数据时，会引起公司的隐私问题。例如，据报道，三星通过 ChatGPT 泄露了自己的秘密，突出了与此类 API 相关的风险。
• 专有的和训练有素的模型是重要的知识产权。对于企业来说，最能从人工智能中受益的问题和数据集往往是敏感和专有的。这就排除了对公共模型的微调，因此业务最好是利用内部部署的模型进行微调，或在企业的防火墙内部署模型。
• 成本问题。目前，LLM 的训练需要数以千计的 GPU 节点，成本高达十亿级美元。模型可能需要平衡不同的权衡，如略微影响模型的性能以减少云的成本。拥有模型使客户可以方便地对其进行微调和再训练。例如，他们可以删除多余的模型参数，以实现特定领域任务的高性能，同时最大限度地降低云成本。

由此，引发出我们这篇文章讨论的方法：FL+LLM，即引入联邦学习来训练大语言模型，从而为企业商业用户提供众多优势，在模型规模和性能、隐私、效率、云计算成本和劳动力成本方面大大增强企业使用大型模型的能力。

1.2 联邦学习回顾

联邦学习（Federated learning，FL）是一种机器学习环境，在这种环境下，多个客户端（如移动设备或整个组织）在中央服务器（如服务提供商）的协调下协同训练一个模型，同时保持训练数据的分散性。FL 体现了集中数据收集和最小化的原则，可以减轻传统的集中式机器学习和数据科学方法带来的许多隐私、安全性风险和成本。因此，FL 是一种有效的高性能计算范式，也被看作是满足数据隐私性要求的分布式训练方法。目前，FL 已经在一些应用场景中成功使用，包括一些消费端设备，例如 GBoard、关键词发现，以及制药、医学研究、金融、制造业等等。此外，也出现了大量 FL 开发工具和平台，例如 Tensorow Federated, LEAF, PaddleFL 和 PySy，进一步推动 FL 的发展。

在联邦学习框架中，中央服务器保存初始化可共享的全局数据。各个客户端（参与者、边缘设备）保存本地数据，并根据本地数据训练本地机器学习模型。客户端根据一定的通信机制向中央服务器传输模型参数等数据（不会传输完整的客户端原始数据），中央服务器汇聚各客户端上载数据后训练构建全局模型，各个客户端在整个联邦学习机制中身份和地位相同。联邦学习有效解决了两方或多方数据使用实体（客户端）在不贡献出数据的情况下的数据共同使用问题，解决了数据孤岛问题。此外，在各个客户端数据特征对齐的前提下，联邦学习的全局模型能够获得与数据集中式存储相同的建模效果。

以下是四个非常重要的联邦学习算法：

1. FedAvg（Federated Averaging）[1]：FedAvg 是最经典的联邦学习算法，由 Google 于 2016 年提出。该算法采用梯度下降的方式对模型参数进行迭代更新。在 FedAvg 中，每个客户端（例如，手机或其他设备）首先在本地用自己的数据训练模型，然后将本地模型权重发送至中央服务器。中央服务器收集所有客户端的权重更新，计算权重的平均值，并将更新后的权重分发回各个客户端。这个过程在多轮中进行，直到达到收敛。

2. FedProx（Federated Proximal）[2]：FedProx 是一种改进的联邦学习算法，用于解决联邦学习中的非独立同分布（non-IID）数据和设备异质性问题。这些问题会导致训练质量下降，尤其是在训练数据不均匀分布的情况下。FedProx 通过在优化目标中添加一个 proximal term 来解决这个问题，这使得模型在更新时更加关注与全局模型权重接近的本地模型权重。这样可以减轻非独立同分布数据和设备异质性带来的负面影响，提高模型性能。

3. SCAFFOLD（Federated Learning via Stochastic Controlled Averaging）[3]：SCAFFOLD 是一种高效的联邦学习算法，主要关注梯度稀疏性和设备间的通信。SCAFFOLD 通过在每轮中维护一个控制变量（control variate）来减少通信量。这个控制变量可以根据全局模型的梯度变化进行调整，从而使每个设备之间的通信更为紧凑。SCAFFOLD 在通信方面更加高效，可以在不牺牲模型性能的情况下减少通信成本。

4. FedNova（Federated Learning via Normalized Averaging）[4]：FedNova 的核心思想是采用归一化平均方法消除目标不一致性，同时保持快速的误差收敛。在每轮迭代中，在各个客户端设备上进行本地训练，并对本地模型参数进行归一化，然后将归一化后的参数发送至中央服务器。中央服务器在收集所有客户端的参数更新后，对其进行归一化平均，从而减少目标不一致性。之后，中央服务器将更新后的模型参数分发回各个客户端，以便在下一轮迭代中继续训练。FedNova 的优点在于，它能够在保持快速误差收敛的同时，消除由于数据异质性导致的目标不一致性。这使得 FedNova 算法在处理非独立同分布数据和设备异质性方面表现出较好的性能。

2、联邦学习结构来训练大语言模型的难点

近年来，LLM 和 FL 的研究都取得了一定的进展。但是，在目前的技术和资源条件下，直接 FL+LLM 还面临一系列的问题。

2.1 数据方面

LLM 的训练依托的是超大规模的数据量。Chinchilla 探讨了训练 LLM 需要多少数据 [9]。2020 年 5 月，OpenAI 展示了他们的 LLM data scaling laws（也称为 Kaplan scaling laws）：300B tokens 可用于训练参数量大小为 175B LLM。2022 年 9 月，DeepMind 为 “数据最优” LLM 找到了新 data scaling laws（也称为 Chinchilla 或 Hoffman scaling laws）：应该使用 1,400B (1.4T) tokens 来训练参数量大小为 70B LLM 最佳。

表 1. 与 Chinchilla 数据优化模型一致所需的数据集大小 [9]

由表 1 中 Chinchilla 给出的数据可以看出，应用联邦学习结构来训练 LLM 势必要将海量的数据量分布在不同的客户端设备中，如果客户端数量较少，那就造成每个客户端中存储、承载的数据量依然较大，对每个客户端的计算性能要求都非常高，这与联邦学习架构的设计初衷是不一致的；如果客户端数量非常庞大，每个客户端中存储的数据量适中，那对于大量客户端之间的协调调用机制、多方计算能力和资源分配不均的资源不均衡问题就会非常突出，不均衡情况下的客户端激励机制也尤为关键，会对联邦学习的中央模型效果、训练效率都有很大的影响。

2.2 客户端设备方面

联邦学习架构下的分布式系统中，单个客户端如何容纳整个 LLM 模型？在已有的联邦学习平台中，特别是一些实用场景下的联邦学习平台，分布式的客户端设备大多为手机、平板等，这些设备根本不可能容纳整个 LLM 模型或为模型训练提供足够的内存资源。如果将这些设备完全排除在训练过程之外，也与联邦学习架构的设计初衷不一致。但在客户端无法容纳整个模型、又不能将客户端完全排除在训练过程之外的情况下，如何对 LLM 进行并行式的训练呢？

微软在一篇文章中对 LLaMA 目前对设备内存的占用量进行了分析 [10]。在模型训练过程中，大部分的内存开销用于 model states (模型参数、优化器的状态、梯度）。剩下的内存开销则来自于 residual states（激活值、临时区缓存和内存碎片）。

以 LLaMA 7B 模型为例，模型参数所需的内存是 (参数量 * 每个参数的内存)：fp32 精度: 7*4 bytes = 28GB；fp16 精度：7*2 bytes = 14GB；int8 精度：7*1 bytes = 7GB；混合精度（fp16/32）训练：存储 fp16 精度 + fp32 精度 = 14GB +28GB = 42GB。

梯度所需的内存（计算方式同模型参数一样，参数量 * 每个参数的内存）：fp32 精度: 7*4 bytes = 28GB；fp16 精度：7*2 bytes= 14GB；int8 精度：7*1 bytes = 7GB；混合精度（fp16/32）训练：只存储 fp16 精度 = 14GB。

优化器所需的内存（参数量 * 每个参数的内存 * 2，以 Adam 为例）：fp32 精度: 7*4 bytes * 2 = 56GB；fp16 精度：7*2 bytes * 2 = 28GB；int8 精度：7*1 bytes * 2 = 14GB；混合精度（fp16/32）训练：存储 fp32 精度 = 56 G。

总的来说，训练 LLaMa 模型总的内存需求为，fp32 精度: 28+28+56 = 112GB；fp16 精度：14+14+28 = 56GB；int8 精度：7+7+14 = 28G；混合精度（fp16/32）训练: 42+14+56 = 112GB。

LLaMA 7B 的规模相当于目前的 LLM 来说非常小了，但即使 LLaMA 7B 的内存消耗量对一些手机、平板来说也是难以满足的，这就对 FL 架构下的客户端设备选择构成了巨大的约束。

2.3 训练算法方面

已有的联邦学习训练算法对于 LLM 来说是低效的。如上所述，联邦学习的工作过程是：“各个客户端根据本地数据训练本地模型。客户端根据一定的通信机制向中央服务器传输模型参数等数据，中央服务器汇聚各客户端上载数据后训练构建全局模型”。由于 LLM 的参数量和模型结构非常庞大，这种传输、汇聚的过程又受到通信条件的影响，所以整个工作过程的效率是非常低的。

3、FL+LLM 已有哪些工作进展

尽管 FL+LLM 还面临诸多问题，但是不少研究人员已经开始探索相关方向的技术可能性并取得了一定的进展。我们在这一章节中介绍一下相关的工作进展情况。

3.1 FedLLM：在专有数据上建立你自己的大型语言模型

FedLLM 是一个 MLOps 支持的训练 pipeline，能够实现在专有数据上建立企业自己的大语言模型 [7]。代码已公开：https://github.com/FedML-AI/FedML/tree/master/python/app/fedllm 。FedLLM 代表了 "LLM 的基础生态系统设计"，而不仅仅是 "LLM 的联邦学习"。

FedLLM 是 FedML.Inc. 推出的 FL+LLM 框架。FedML, Inc. (https://FedML.ai ) 是由华人主导的国际化团队，他们起源于美国南加州大学 (University of Southern California)，是全球范围内研究该技术的早期机构之一。过去几年，FedML 起初作为博士生主导的科研开源项目，服务于多个科研基金，辅助所在实验室发表了 50 多篇相关顶级论文。2023 年，FedML 宣布完成 600 万美元种子轮和 Pre-Seed 轮融资，Camford Capital 领投，Plug and Play Ventures、AimTop Ventures、Acequia Capital、LDV Partners 等投资者参投。

图 2. FedLLM 的训练过程

如图 2，FedLLM 实现了数据协作、计算协作和模型协作，并支持在集中式和地理分布式的 GPU 集群上进行训练，以及数据孤岛的联邦学习。FedLLM 与 HuggingFace 和 DeepSpeed 等流行的 LLM 库兼容，并旨在提高效率和安全 / 隐私。

• 数据协作 — 在特定领域的专有数据上实现 LLM 的训练。通用的 LLMs ChatGPT 和 GPT-4 是在人类发布和注释的大量文本上训练的。在许多垂直领域（如医疗保健、金融科技、法律和汽车行业），ChatGPT 可能无法很好地发挥作用。使企业能够在其专有数据上训练他们的模型，可以在保护隐私的同时实现更好的性能。在某些情况下，当数据集分散在各个数仓中时，FedML 联邦学习训练管道 "Train on the Edge" 将以安全和可扩展的方式处理。
• 计算协作 — 利用分散的计算资源。由于计算能力的高投资，只有少数大型科技公司（如微软、谷歌）才能负担得起 LLM 预训练。处于不同垂直领域的企业无法负担花费数十亿美元的数千个 GPU 节点。充分利用 GPU 资源的一个更具成本效益的方法是建立一个跨组织的共享机制。FedML 通过 "Train on the cloud" 平台实现了这一点，它可以将训练工作安排到地理分布的 CPU/GPU 网络上。这种计算方面的合作可以减少在一个组织中购买大量 GPU 节点的财务负担。
• 模型协作 — 以联合的方式服务于模型。服务大型基础模型也是一种挑战。与其他 MLOps 相比，FedML 率先提出了通过地理分布的云计算资源进行联邦模型推理的想法。当推理请求被发送到推理端点时，主节点将把请求路由到由 GPU 供应商托管的分散的边缘节点，他们可以共享 GPU 的空闲时间。这样的模式服务平台可以提供更好的服务可靠性和更便宜的云计算成本。

对于特定的孤岛式 GPU 集群，FedLLM 利用现有的开源 LLM 和流行框架进行本地训练：

• 模型定义和预训练的权重来自 EleutherAI 的 pythia 的 2.8B、7B 和 12B 版本。将 Hugging Face 的 transformer 作为参考实现（https://github.com/huggingface/peft?ref=blog.fedml.ai）。FedLLM 也是一个模型无关的框架，企业和开发者可以插入任何 LLM。
• FedLLM 支持参数高效的训练方法，如 LoRA。其参考实现来自 HuggingFace 的 peft（https://huggingface.co/docs/transformers/index?ref=blog.fedml.ai）。
• 对于一个独立的训练器，其训练 / 评估代码是基于 Transformers 库（https://huggingface.co/docs/transformers/index?ref=blog.fedml.ai）中的 Trainer。
• 特定 GPU 集群上的分布式训练框架由 DeepSpeed（https://www.deepspeed.ai/?ref=blog.fedml.ai）处理。还启用 Zero3 来减少单个 GPU 的内存成本。通过加强 FedML 和 DeepSpeed 之间的兼容性，FedLLM 可以在不同的物理空间分布式集群中运行训练作业。

FedML 也给出了一个应用 FedLLM 的实例。如图 3，本地数据保持在本地，只有模型参数或权重数据从中央服务器流向和流回。这个特殊的实例假设了一个聊天类型的应用，它结合了本地数据的训练，以及中央服务器通过利用其他设备的训练而建立的微调训练的好处。聊天应用程序只是一个例子，可以被其他利用 LLM 模型的应用程序所取代。

图 3. FedLLM 应用实例

FedLLM 从一定程度上能够实现利用联邦学习训练 LLM。但是，回顾我们在第二章中讨论的，利用联邦学习架构训练 LLM 还面临很多问题，从公开的技术内容来看，FedLLM 并没有能够解决这些问题。FedLLM 架构中假设的计算节点都是具备一定计算能力的 GPU 节点，FedLLM 能够解决的问题是 “不用大量购置 GPU”。但真实状况下的联邦学习，客户端节点很多是手机、平板，FedLLM 并不能适配这些情况。同样的问题也存在于客户端节点数据量的情况，即 FedLLM 假设客户端是足够容纳训练本地 LLM 所需要的数据量的。最后，FedLLM 并没有讨论什么样的聚合算法适合 FL for LLM，也没有讨论是否需要改进客户端训练 LLM 的算法，这对于真正推广应用 FL+LLM 还是需要进一步细化解决的。

3.2 DeepSpeed：通过提高规模、速度、成本和可用性来推进大型模型训练，释放出训练 1000 亿个参数模型的能力

LLM 在训练时往往需要大量内存来存储中间激活、权重等参数，百亿模型甚至无法在单个 GPU 上进行训练，使得模型训练在某些情况下非常低效和不可能。前期工作已经有研究人员专门聚焦在多节点分布式训练问题中，主要采用数据并行（模型的不同实例在不同的 GPU 和不同批数据上运行）和模型并行（将模型拆分到多个 GPU 上进行训练）两种方式。

DeepSpeed 是微软开发的一种深度学习训练优化工具，旨在通过分布式训练和混合精度技术来提高训练速度和节省资源。它是一个开源的 Python 库，可以在多个平台上运行。与传统的深度学习框架，例如 TensorFlow、Pytorch、Keras 等相比，DeepSpeed 通过将模型参数拆散分布到各个 GPU 上，以实现大模型的计算，使得可以用更少的 GPU 训练更大的模型，而且不受限于显存。

数据并行会在所有 worker 之间进行模型和优化器的复制，因此显存效率不高。DeepSpeed 开发了 ZeRO ，它是一系列用于提高数据并行的显存效率的优化器。针对模型并行，DeepSpeed 借用了英伟达的 Megatron-LM 来为基于 Transformer 的语言模型提供大规模模型并行功能。

具体来说，DeepSpeed (https://www.deepspeed.ai) 能够通过一次点击实现类似 ChatGPT 的模型训练，与 SOTA RLHF 系统相比，提供了 15 倍的速度提升，并在所有规模上实现了前所未有的成本降低。DeepSpeed 有三个创新点：

• DeepSpeed-Training：DeepSpeed 提供了一系列的系统创新，这些创新使得大规模的深度学习训练变得有效和高效，极大地改善了使用的便利性，并且在可能的规模方面重新定义了深度学习训练领域。这些创新，例如 ZeRO、3D-Parallelism、DeepSpeed-MoE、ZeRO-Infinity 等，都属于 DeepSpeed-Training 的核心。
• DeepSpeed inference：DeepSpeed inference 是一个用于 Transformer 模型推理的全面系统解决方案。DeepSpeed inference 包括：(1) 一个多 GPU 推理解决方案，以最大限度地减少延迟，同时最大限度地提高密集和稀疏 Transformer 模型的吞吐量，当它们适合聚合 GPU 内存时；(2) 一个异构推理解决方案，除了 GPU 内存和计算外，还利用 CPU 和 NVMe 内存，以实现不适合聚合 GPU 内存的大模型的高推理吞吐量。
• DeepSpeed-Compression：为了进一步提高推理效率，DeepSpeed 为研究人员和从业人员提供了易于使用和灵活组成的压缩技术，以压缩他们的模型，同时提供更快的速度、更小的模型尺寸，并大大降低压缩成本。此外，SoTA 在压缩方面的创新，如 ZeroQuant 和 XTC，都包含在 DeepSpeed-Compression。

具体来说，DeepSpeed Inference 的两个核心组件介绍如下：

• DeepSpeed Transformer：DeepSpeed Transformer 是一个仅适用于 GPU 的解决方案，旨在最小化延迟，同时最大限度地提高密集和稀疏 Transformer 模型的吞吐量。它为各种规模的 Transformer 模型实现了最先进的延迟和吞吐量，并支持在单个 GPU 上运行或扩展到数百个 GPU 以推理数百万参数的模型。DeepSpeed Transformer 解决方案是一个三层的系统架构，包括：i) 单个 GPU 转化器内核，在小批规模时针对内存带宽利用率进行优化，在大批规模时针对高吞吐量进行优化；ii) 多个 GPU 密集 Transformer 层，利用张量切割和推理优化的管道并行性在 GPU 上扩展密集 Transformer 模型。海量 GPU 规模的稀疏 Transformer 层，旨在利用并行技术和通信优化策略的组合将 MoE Transformer 层扩展到数百个 GPU，同时还利用优化的稀疏内核将单 GPU 稀疏计算的开销降到最低。通过采取这种分层方法，每一层都解决了延迟挑战的一个独特方面：批大小、密集模型的扩展和稀疏模型的扩展，但又相互兼容并建立在彼此之上，DeepSpeed Transformer 是一个综合系统，尽管在批大小、模型规模和模型特征方面存在异质性，但它能够以前所未有的规模为密集和稀疏 Transformer 模型实现最优的延迟和吞吐量。
• ZeRO-Inference：ZeRO-Inference 是一个基于异构 GPU+CPU+NVMe 的解决方案，是一种用于大规模分布式深度学习的新型内存优化技术。通过用最少的 GPU 资源实现大规模的模型推理，来解决内存的挑战。与 DeepSpeed Transformer 相比，对于那些对延迟不太敏感但资源有限的应用，只要有足够的 CPU 或 NVMe 内存来存储模型参数，ZeRO-Inference 就可以在单个或多个 GPU 上推断具有数千亿个参数的模型。

图 4. DeepSpeed 压缩库

DeepSpeed Compression 提出了一个无缝 pipeline 来解决压缩合成性的挑战，如图 4 所示。DeepSpeed Compression 的核心部分是一个叫做 compression composer 的组件，它包括几个重要的功能：

1. 它提供了多种尖端的压缩方法，包括极端量化、头 / 行 / 通道修剪和知识提炼，可以有效地减少模型大小和推理成本。随着不断地整合更多的最先进的压缩方法，这个列表还会扩大。

2. 它提供了一个易于使用的 API，自动处理组装不同压缩技术的复杂问题，以提供多种压缩方法的复合效益。例如，XTC 需要组成轻量级的减层、二值化和知识蒸馏。然而，将它们组合在一起并不困难。有了 compression composer，应用极端压缩就像添加两个新的 API 调用来启用压缩和清理压缩模型一样简单。

3. 它是以模块化的方式设计的，这样用户就可以很容易地添加新的压缩方案。例如，额外的压缩方法可以通过自定义压缩层添加，通过向 compression composer 注册，新方法可以与已经由 composer 管理的现有方法组成。

4. 它与 DeepSpeed 库无缝衔接。这有两个好处。首先，DeepSpeed 压缩可以通过 JSON 文件以与 DeepSpeed 训练和推理相同的方式进行指定和启用，其中启用不同的压缩技术组合只需要在 JSON 文件中修改几行。其次，一旦配置了压缩方案，压缩 composer 就会自动修改模型层和训练以启用压缩过程，不需要用户对模型结构或训练程序进行额外的修改。

3.3 相关算法基础

为了应对 LLM 大规模预训练的成本问题，从算法层面，研究人员提出了 LoRA（ LOW-RANK ADAPTATION ），即大语言模型的低秩适应方法，一种轻量级训练大语言模型的方法。通过冻结预训练的模型权重，并将可训练的秩分解矩阵注入到 Transformer 架构的每一层，极大地减少了下游任务的可训练参数的数量，有效提升预训练模型在下游任务上的 finetune 效率。

项目地址：https://github.com/microsoft/LoRA [6]

为解决微调大语言模型到不同领域和任务的挑战，已有多种方案，比如部分微调、使用 adapters 和 prompting。但这些方法存在如下问题：

• Adapters 引入额外的推理延迟 (由于增加了模型层数)
• Prefix-Tuning 难于训练，且预留给 prompt 的序列挤占了下游任务的输入序列空间，影响模型性能

Lora 提出的思路是：虽然模型的参数众多，但其实模型主要依赖低秩维度的内容 (low intrinsic dimension)。假设在适应过程中，权重的更新也具有较低的 "内在等级（intrinsic rank）"。对于一个预先训练好的权重矩阵 W_0∈R^d×k，通过用低秩分解代表后者来约束其更新：

,

,

在训练期间，冻结 W_0，不接受梯度更新，而 A 和 B 包含可训练参数。W_0 和∆W = BA 都与相同的输入相乘，它们各自的输出向量按坐标相加。对于 h=W_0x，修改后的正向传递得到了：

图 5. Reparameterized 只训练 A 和 B

如图 5，对 A 使用随机高斯初始化，对 B 使用零初始化，所以∆W = BA 在训练开始时是零。然后，用 α/r 来调整∆Wx，其中 α 是 r 中的一个常数。当用 Adam 进行优化时，如果适当地调整初始化，那么调整 α 与调整学习率大致相同。因此，只需将 α 设置为尝试的第一个 r，而不对其进行调整。这种比例有助于减少在改变 r 时重新调整超参数的需要。

图 6. 常规微调和 LoRA 图示

如图 6，r 是超参数，指定用于自适应的低秩矩阵的秩。r 越小，低秩矩阵越简单，在自适应过程中需要学习的参数越少，训练就更快，计算需求会相应的减少。然而，r 变小的弊端是，低秩矩阵捕获任务特定信息的能力降低。这可能导致较低的自适应质量，并且与较高的 r 相比，模型在新任务上可能表现不佳。总之，在 LoRA 中确定 r 的取值，需要在模型复杂性、自适应能力和拟合不足或拟合过度的风险之间进行权衡。因此，重要的是用不同的 r 值进行实验，以找到正确的平衡，从而在新任务中满足所需的性能。

矩阵秩 (r) 可以非常低，比如对于 GPT-3 175B 模型，rank-1 或 rank-2 就能基本对齐原始 rank-12288 的效果：

• 与 Adam 微调的 GPT-3 175B 相比，LoRA 可训练参数数量减少了 1 万倍，GPU 内存需求减少了 3 倍
• 在 RoBERTa、DeBERTa、GPT-2 和 GPT-3 等大语言模型上，LoRA 在模型质量方面的表现与微调相当或更好，尽管它具有更少的可训练参数、更高的训练吞吐量，并且与适配器不同，没有额外的推断延迟

3.4 相关硬件基础

最后，我们来看看硬件方面的工作。

Grace Hopper 是第九代英伟达数据中心 GPU，与前几代英伟达安培 GPU 相比，旨在为大规模 AI 和 HPC 应用提供数量级的改进。线程块集群和线程块重新配置改善了空间和时间上的数据定位，再加上新的异步执行引擎，使应用程序能够始终保持所有单元的忙碌。NVIDIA Grace Hopper 通过 NVIDIA NVLink C2C 将 NVIDIA Grace CPU 和 NVIDIA Hopper GPU 融合到一个超级芯片中，NVLink C2C 是一个总带宽为 900 GB/s 的芯片间互连。NVLink C2C 内存一致性使得 Grace CPU 超级芯片和 Grace Hopper 超级芯片都可以通过统一的编程模型进行编程。

图 7. 英伟达 Grace Hopper 超级芯片逻辑概览

2023 年 6 月，英伟达宣布其新的 GH200 Grace Hopper "超级芯片"— 专门为大规模人工智能应用创建的 CPU 和 GPU 组合 — 已经进入全面生产。它有 528 个 GPU 张量核心，支持高达 480GB 的 CPU 内存和 96GB 的 GPU 内存，并拥有高达 4TB / 秒的 GPU 内存带宽。

图 8. 英伟达的 GH200 "Grace Hopper"AI 超级芯片 [5]

GH200 以 "Hopper" 为基础，与 Nvidia 的 "Grace"CPU 平台（两者均以计算机先驱 Grace Hopper 命名）相结合，通过 Nvidia 的 NVLink 芯片对芯片（C2C）互连技术将其纳入一个芯片。Nvidia 预计，这一组合将极大地加速人工智能和机器学习应用的训练（创建模型）和推理（运行模型）。

"生成式人工智能正在迅速改变企业，在医疗、金融、商业服务和更多行业释放新的机会并加速发现，"Nvidia 加速计算副总裁 Ian Buck 在一份新闻稿中说。"随着 Grace Hopper 超级芯片的全面生产，全球制造商将很快提供企业所需的加速基础设施，以建立和部署利用其独特的专有数据的生成式人工智能应用程序。" 据该公司称，GH200 的主要特点包括一个新的 900GB/s 相干（共享）内存接口，比 PCIe Gen5 快 7 倍。GH200 还为 GPU 提供了 30 倍的系统内存总带宽。此外，GH200 可以运行所有 Nvidia 软件平台，包括 Nvidia HPC SDK、Nvidia AI 和 Nvidia Omniverse。Nvidia 还宣布，它将把这种 CPU/GPU 组合芯片构建成一个新的超级计算机，称为 DGX GH200，它可以利用 256 个 GH200 芯片的综合能力，作为一个单一的 GPU 来执行，提供 1 exaflop 的性能和 144 兆字节的共享内存，比上一代 Nvidia DGX A100 的内存多近 500 倍。DGX GH200 将能够为生成语言应用、推荐系统和数据分析训练巨大的下一代人工智能模型。Nvidia 目前并未公开 DGX GH200 的售价，但如果根据去年出货的 DGX H100 作为参考，一台搭载 8 组 H100 GPU 的 8U GPU 服务器机柜约为 20 万美元，考虑到 DGX GH200 拥有最多 256 个 Grace Hopper，其价格可能会高于这个范围。根据 Anandtech 的报道，一台 DGX GH200 计算机的价格 "很容易达到 8 位数（美元）的水平"。

由于 Nvidia 和 Cerebras 等供应商的持续硬件进步，高端云人工智能模型可能会随着时间的推移继续变得更有能力，处理更多的数据，并且比以前做得更快。根据内部测试，在处理需要大量内存的 AI 工作负载方面，DGX GH200 系统显示出比 DGX H100 更好表现，平均可获得 2-6 倍不等的性能提升，例如，在使用 1TB 的内存容量用于 GPT3 模型训练时，DGX GH200 的处理速度快上 1 倍，若以 40TB 内存处理深度学习推荐模型（DLRM）时，速度则有多达 4 倍的提升，甚至在图神经网络处理方面速度更快上许多，足足有 5 倍之多。

单纯从 GH200 的角度看，能够在一定程度上支撑对 LLM 的分布式训练，使 FL+LLM 变得更加可行。但是，回到文初我们的讨论，FL 的初衷是利用大量分散的一般性设备分布式训练一个中央模型，一方面可以有效利用分散的客户端资源，另一方面也满足各个客户端的数据隐私需要。要求这些客户端都装有 GH200 显然是不现实的，这样的 GPU 的成本消耗与 FL 的初衷并不一致。

此外，也有一些研究人员认为 GH200 自身也不会对大模型的应用推广有太大影响 [11]。他们分析，算力层面，单颗 GH 芯片和 H100 芯片 FP8 浮点算力没有差异。内部测试中 GH200 更好的表现是由于其存储。DGX GH200 内部 GPU 和 CPU 的连接方式与 DGX H100 不同，其高速触达的存储容量可以大幅提升。然而，集群性能的提升主要有算力本身、网络、存储三大要素。因此，传统 GPT3、GPT4 等主流大模型下，DGX H100 集群（NVLink 联网）与 DGX GH200 系统并不会有明显差异。

关于 H100，已有报道分析了其大模型训练的性能 [12]。两项 MLPerf 基准评测的最新数据中，英伟达 H100 芯片组在人工智能算力表现的测试中，刷新了所有组别的纪录，也是唯一一个能够跑完所有测试的硬件平台。英伟达和 GPU 云算力平台 CoreWeave 提交的数据为这项测试设定了业界标准。在 896 个英特尔至强 8462Y + 处理器和 3584 个英伟达 H100 芯片的联合作用下，仅仅花了 10.94 分钟就完成了基于 GPT-3 的大语言模型训练任务。在 BERT-Large 模型训练中，H100 和 CoreWeave 将数据刷到了极端的 0.13 分钟，在 64 卡的情况下，测试数据也达到了 0.89 分钟。

4、后续发展讨论

我们在这篇文章中讨论了与 FL+LLM 相关的一些方法，包括算法方面的改进、硬件方面的研发，以及分布式训练和联邦学习的架构。其中，FedLLM 应当是最符合 FL+LLM 的一篇工作，尽管其距离实用、技术完备的 FL+LLM 还有很大差距。

DeepSpeed 是通过分布式训练和混合精度技术来提高训练速度和节省资源的模型训练架构，其面向的是分布式训练的应用场景，一般来说，分布式学习会将训练数据均匀分布在不同节点，节点大多是专用机房中的计算节点，且节点通常处于统一地理位置，节点间通信条件较好。而对于 FL 来说，每个计算节点拥有的数据量与设备自身有关，很难保证不同计算节点拥有相近的数据量。此外，FL 的计算节点可能是手机、平板等，一般与中央服务器处于远程连接状态，连接稳定性和通信代价的情况都比较差。所以目前 DeepSpeed 的工作直接应用到 FL+LLM 中还有很大难度。

LoRA 是目前非常流行的微调方法，更多优秀的微调方法也在不断提出。算法上的改进，给 FL 各个客户端节点基于本地数据训练模型带来了可能。而硬件目前的发展我们在前面也讨论过，硬件自身的性能在不断提高，但对于 FL 这个应用场景来说，令手机、平板具备这样的硬件条件难度非常大。

由对目前与 FL+LLM 相关的一些工作的分析，我们感觉联邦学习支撑 LLM 还有很多需要解决的问题。包括 FedLLM，实际上也没有谈及如何应对存储容量小、处理性能差、网络条件不佳的客户端设备，而这些，恰好是 FL 最典型的应用场景。我们期待有更多的研究人员可以关注 FL+LLM 问题，给出更多的技术细节和可落地的解决方案。

作者介绍

Jiying，工学博士，毕业于北京交通大学，曾分别于香港中文大学和香港科技大学担任助理研究员和研究助理，现从事电子政务领域信息化新技术研究工作。主要研究方向为模式识别、计算机视觉，爱好科研，希望能保持学习、不断进步。

参考引用的文献：

[1] McMahan, Brendan, et al. "Communication-efficient learning of deep networks from decentralized data."Artificial intelligence and statistics. PMLR, 2017.

[2] Li, Tian, et al. "Federated optimization in heterogeneous networks."Proceedings of Machine learning and systems2 (2020): 429-450.

[3] Karimireddy, Sai Praneeth, et al. "Scaffold: Stochastic controlled averaging for federated learning."International Conference on Machine Learning. PMLR, 2020.

[4] Wang, Jianyu, et al. "Tackling the objective inconsistency problem in heterogeneous federated optimization."Advances in neural information processing systems 33 (2020): 7611-7623.

[5] https://arstechnica.com/information-technology/2023/06/nvidias-new-ai-superchip-combines-cpu-and-gpu-to-train-monster-ai-systems/

[6] LoRA: Low-Rank Adaptation of Large Language Models， https://arxiv.org/abs/2106.09685

[7] Releasing FedLLM: Build Your Own Large Language Models on Proprietary... (https://blog.fedml.ai/releasing-fedllm-build-your-own-large-language-models-on-proprietary-data-using-the-fedml-platform/)

[8] Wayne Xin Zhao, Kun Zhou, etc., A Survey of Large Language Models, https://arxiv.org/abs/2303.18223

[9] https://lifearchitect.ai/chinchilla/

[10] Samyam Rajbhandari, Jeff Rasley, etc., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, https://arxiv.org/pdf/1910.02054.pdf%22%20%5Ct%20%22_blank

[11] 国泰君安证券，研究报告《AI 超算浑然一体，光互联地位显著提升》

[12] 英伟达 H100 霸榜权威 AI 性能测试，11 分钟搞定基于 GPT-3 的大模型训练，https://m.cnbeta.com.tw/view/1367739.htm