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

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 有三个创新点：