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

“隐语”是开源的可信隐衷计算框架，内置 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
本文作者：隐语技术团队 黄炜

一、背景

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

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

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

二、先验常识

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