「隐语小课」针对大模型！护航数据安全，一文详解如何实现参数可控的分布式训练

“隐语”是开源的可信隐私计算框架，内置 MPC、TEE、同态等多种密态计算虚拟设备供灵活选择，提供丰富的联邦学习算法和差分隐私机制。

开源项目：
https://github.com/secretflow
https://gitee.com/secretflow

导语：聚焦大模型隐私保护前沿技术探索，本文将对隐语团队发表在 ICASSP 2024中的论文——A FAST, PERFORMANT, SECURE DISTRIBUTED TRAINING FRAMEWORK FOR LLM，进行解读。
🔍 原文链接：https://arxiv.org/abs/2401.09796
本文作者：隐语技术团队黄炜

一、背景

随着大模型的兴起，学术界与工业界都在进行大模型在垂直领域的应用探索。当大模型应用在特定行业时，需要高质量专业数据用于微调。然而每个机构有限的专业数据限制了模型的精度。因此，分布式 (联邦) 大模型成为了备受关注的一个方向。然而，各个机构的数据中往往存在着很多的用户信息以及独有的来源，使得其成为极为重要的的财产。鉴于大模型对数据高效的学习以及记忆能力，导致分布式大模型在训练阶段存在一些安全问题：

恶意的服务端会窃取模型的参数并从中推断客户端的数据。
恶意的客户端也会通过模型参数以及前向 embedding 窃取其余客户端的数据。
微调的参数被各方持有，可能会引起参数的不可控和滥用。

基于上述安全性考虑，我们提出了新的场景，即参数可控的分布式训练。即在模型微调过程中任何参与方都无法拿到微调的参数并且能保证各个参与方数据的安全性。

二、先验知识

2.1 联邦学习

联邦学习是一种保护数据隐私的多方联合建模的训练方法，由 Google 在2016年首次提出。联邦学习借鉴了参数服务器的思路，参与训练的成员分为服务器端和客户端，客户端拥有数据集，服务器负责进行模型聚合。具体训练过程如下：客户端利用自身数据集进行训练，更新本地模型；客户端将模型上传至服务器端，服务器端进行模型融合，得到全局模型；客户端从服务器获得全局模型，更新本地模型，完成一次训练过程。联邦学习通过共享模型而非数据的方式，防止隐私数据的泄漏。

2.2 可信执行环境

可信执行环境 (TEE) 是一种具有运算和储存功能，能提供安全性和完整性保护的独立处理环境。其基本思想是：在硬件中为敏感数据单独分配一块隔离的内存，所有敏感数据的计算均在这块内存中进行，并且除了经过授权的接口外，硬件中的其他部分不能访问这块隔离的内存中的信息。以此来实现敏感数据的隐私计算。流行的 TEE 实现包括英特尔 SGX [1] 、AMD SEV[3]、英特尔 TDX[2] 和 TrustZone[4]。在本文中，我们将 TEE 视为潜在敌手主机设备（包括 GPU）上的一个安全区域。这意味着 TEEs 内部的数据、代码及整个计算过程都是安全的。

三、算法模型

3.1 方法1：小内存（消费级）TEE 保护的LLM拆分

在分布式场景中，假设有 K 个客户端，Let D_k = {(x^m,y^m)}(m=1)^N ,i=1,…,K 指代第 K 个客户端的训练集。在联邦训练可见如下图1(a)，Method1在服务端和客户端之前只通信 LoRA 以及 P-Tuning v2 的参数。我们使用 W_k 指代第k个客户端需要更新的参数。然后在第t轮，全局模型参数 W^' 可以表示为
，其中是全局联合数据的总数量。为了解决服务端模型参数和数据的泄漏问题，我们在服务端部署 TEE(SGX) 并在其中做参数聚合。在训练过程中，我们固定除了 Wk 以外的参数，以学习模型参数 θ，并使用以下目标函数：

其中 L 是第 k 个客户端的损失函数。由于当前客户端会从本地模型的参数和前向 embedding 中恶意窃取其余客户端的数据。为了避免上述问题，我们在每个客户端皆部署了一个 TEE(SGX)，并且把模型的 LoRA 以及 P-Tuning v2 放入该设备。但是 TEE 需要与 GPU 中的大模型交互，由于 TEE 的输出 embedding 中携带大量的用户信息，这也会带来较大的隐私风险。我们通过 one-time pad (OTP) [5] 加密 GPU 和 TEE 之间的传输特征来解决上述问题。由于 OTP 只能对线性运算做加密以及解密，所以除了微调结构外，大模型的非线性层以及运算都部署在 TEE, 包括 layer norm、activation function 以及softmax。具体的划分可见如下图1。

Feature Encryption. 对于大模型的线性层h(.), letE 作为 TEE 的明文输出。然后我们生成一个随机掩码r, 并且通过如下公式加密特征：

Feature Decryption. GPU 接收，计算
,并且返回结果到 TEE。TEE 解密结果如下：

由于 Q∙K 涉及到矩阵乘法，考虑到此计算在TEE运算较为费时，所以我们在 GPU运行此计算。矩阵乘的解密公式如下：

上述公式中，Q_E 和 K_E 指代 Q 和 K 的明文矩阵，Q_En 和 K_En 指代 Q 和 K 的明文矩阵, r 是对应的掩码。Method1 将分布式大模型切分，将敏感部分结构部署在 TEE，并通过 OTP 保护 TEE 与 GPU 的传输，以达到模型参数和数据的保护效果。

3.2 方法2：大内存 TEE 保护的 LLM 拆分

Method1 解决了安全问题，但是其客户端在每个 batch 训练时 GPU 和 TEE 会产生较多传输次数。由于大模型的前向 embedding 纬度较大，会加大传输时延，从而造成训练时间过长的问题。其次，在加密以及解密的过程中，会造成模型精度下降。为了缓解上述问题，我们在 Method2 中提出了一种切分微调的方法，并将其运用到分布式大模型。模型结构和微调方法可见下图2。

Method2 首先按层将模型切分，分为两部分，let M_f 指代前半部分模型，M_l 指代后半部分模型，E_f 指代前半部分模型最后一层的输出特征。每个客户端我们都部署一个 M_f，服务端部署一个 TEE(Intel TDX) 并将 M_l 部署在该 TEE 中。在分布式训练时，每个客户端冻结 M_f 的参数，并把每一条数据产生的E_f保存。当每个客户端搜集完所有数据的 embedding 后加密一次上传到服务端的TEE。服务端接收到数据后做解密，得到未加密前的 E_f，然后每个 E_f 作为 M_l 的输入用于微调模型。微调策略如图2所示。

特别地，我们对 M_l 的每个部分中的 QKV 线性层和 dense 线性层进行微调，基于这样的考虑，即模量值较大的参数在模型中扮演的角色比模量值较小的参数更为重要，我们提出了 Sparsification Parameter Fine-tuning (SPF) 策略。具体来说，我们用 W_all 来表示线性层的权重，并计算 W_all 的第二维上每个 head 的 L1 范数。我们根据 L1 范数的大小从大到小排序 head。之后，我们将这些 head 按一定比例重组为权重，记为W_train（W_freeze表示其余head的权重）。在微调过程中，我们冻结 W_freeze，只更新 W_train。我们使用 X_train 和 X_freeze 来表示线性层的输入。SPF的前向公式可以表示如下：

SPF方法缓解了需要大量微调参数来完全微调 ML 中的 QKV 线性和 dense 线性层的问题。为了进一步减少微调参数的数量，我们选择冻结 MLP 参数。为了确保性能，我们对接到 MLP 层的 LoRA 进行微调。Method2 每个客户端之前没有交互，所以不用部署 TEE。该方案不但能防止客户端向服务端传输的 Embedding 被窃取，而且反向传播时梯度不需要回传到客户端，避免信息进一步泄漏。因为 M_f 参数冻结且不用更新，故每个客户端可以离线生成 E_f 并一次上传。由于分布式训练结束后，每个客户端并不能得到更新的模型参数，所以该方案由服务端向外提供服务，接受用户的 query 查询。

四、实验设置和结构

4.1 实验配置

我们选取了五个医疗数据集作为此实验的行业数据，分别是 CHIP-CTC、KUAKE-IR、KUAKE-QIC、KUAKE-QQR和KUAKE-QTR。我们选择 ChatGLM-6B 作为基础的大型语言模型。对于 Method2，我们设置分割位置为第24层。对于 QKV 线性和 dense 线性层，我们分别选择了比例（12.5%，25%）和（50%，62.5%）。据我们所知，我们是首次提出在大型语言模型（LLM）分布式训练期间，防止服务器和客户端同时恶意盗取模型参数和数据，以及任何参与方都无法获取微调的参数。我们将其与无任何安全措施的明文联邦 LLM（FL-LLM）以及通过 TEE 保护全模型（SWMT）进行比较。表3 展示了我们提出的两种方案与其他方案之间的比较。

4.2 方法一和方法二的实验结果

表1 显示了 Method1 和 Method2 在五个数据集上的准确率以及需要微调的参数数量。从表中我们可以看出，Method1 在五个数据集上的平均准确率略低于 FL-LLM。其次，Method2 在五个数据集上的准确率显著高于其他方法，这表明Method2的分割微调策略能有效提高下游任务的准确率。然而，Method2 中微调的参数数量大约是 Method1 的5倍。表4 显示了不同方法的训练和推理时间。我们可以得出结论，SWMT 的训练和推理时间最长。由于 TEE 中只有最后 m 层的操作，Method2 有更短的训练时间。尽管 FL-LLM 的延迟很小，但它不能保证模型参数和数据的安全性。

4.3 方法二的消融

实验表2展示了 Method2 的模型准确率和按不同层次和选定比例分段的参数数量。如我们所见，服务器部署的层数越多或选定的百分比越高，需要微调的参数数量就越大，这反过来增加了训练和推理的延迟，但带来了一些性能提升。

五、参考文献

[1] Frank McKeen, Ilya Alexandrovich, Ittai Anati, Dror Caspi, Simon Johnson, Rebekah Leslie-Hurd, and Carlos Rozas, “Intel. software guard extensions (intel.sgx) support for dynamic memory management insidean enclave,” in Proceedings of the Hardware and Architectural Support for Security and Privacy 2016, pp. 1–9.2016.
[2] “Intel. intel trust domain extensions,” https://www.intel.com/content/www/us/en/developer/to ols/trust-domain-extensions/documentation.html.
[3] David Kaplan, Jeremy Powell, and Tom Woller, “Amd memory encryption,” White paper, p. 13, 2016.
[4] Tiago Alves, “Trustzone: Integrated hardware and software security,” Information Quarterly, vol. 3, pp. 18–24, 2004.
[5] Florian Tramer and Dan Boneh, “Slalom: Fast, verifiable and private execution of neural networks in trusted hardware,” arXiv preprint arXiv:1806.03287, 2018.

🏠 隐语社区：
https://github.com/secretflow
https://gitee.com/secretflow
https://www.secretflow.org.cn（官网）

👇欢迎关注：
公众号：隐语的小剧场
B站：隐语secretflow
邮箱：secretflow-contact@service.alipay.com