[论文解读] A Fast, Performant, Secure Distributed Training Framework For Large Language Model

Wei Huang, Yinggui Wang+4 more|arXiv (Cornell University)|Jan 18, 2024

Topic Modeling被引用 4

一句话总结

该论文提出了一种安全、高性能的大型语言模型（LLM）分布式训练框架，采用模型切分、可信执行环境（TEE）和轻量级加密技术，防止服务器端和客户端的模型/数据泄露。通过在客户端和服务器端均部署TEE，并将微调过程按层拆分——尤其将后期层置于服务器端TEE中——该方法在实现强安全性的同时，仅造成极小的精度损失，在准确率和效率方面优于基线方法。

ABSTRACT

The distributed (federated) LLM is an important method for co-training the domain-specific LLM using siloed data. However, maliciously stealing model parameters and data from the server or client side has become an urgent problem to be solved. In this paper, we propose a secure distributed LLM based on model slicing. In this case, we deploy the Trusted Execution Environment (TEE) on both the client and server side, and put the fine-tuned structure (LoRA or embedding of P-tuning v2) into the TEE. Then, secure communication is executed in the TEE and general environments through lightweight encryption. In order to further reduce the equipment cost as well as increase the model performance and accuracy, we propose a split fine-tuning scheme. In particular, we split the LLM by layers and place the latter layers in a server-side TEE (the client does not need a TEE). We then combine the proposed Sparsification Parameter Fine-tuning (SPF) with the LoRA part to improve the accuracy of the downstream task. Numerous experiments have shown that our method guarantees accuracy while maintaining security.

研究动机与目标

• 为解决分布式（联邦）LLM训练中模型参数和数据泄露的关键安全挑战，特别是来自恶意服务器和客户端的泄露风险。
• 设计一种安全、高效且可扩展的训练框架，在保护敏感数据和模型权重的同时保持高模型准确率。
• 通过支持消费级（小内存）和工业级（大内存）TEE解决方案，实现跨多样化硬件的实用化部署。
• 通过一种新颖的分层微调策略，最小化TEE中更新的参数数量，从而降低训练开销和设备成本。

提出的方法

• 该框架使用模型切分技术将LLM进行分割，将后期层置于服务器端TEE中，而前期层则保留在客户端的一般计算环境中。
• LoRA和P-tuning v2参数在客户端和服务器端的TEE中存储并计算，以确保微调权重和嵌入表示的机密性。
• 通过一次性密码本（OTP）加密实现TEE与外部环境之间的安全通信，用于梯度和中间输出的保护。
• 引入一种分层微调策略，仅在服务器端TEE中微调部分层（例如最后4层），从而减少需在TEE中更新的参数面和成本。
• 该方法支持两种TEE配置：Intel SGX（小内存，无GPU）和Intel TDX/SGX（大内存，无GPU），并为每种配置设计了定制化方案。
• 将稀疏化参数微调（SPF）与LoRA结合，进一步提升下游任务的准确率，且不增加模型大小。

实验结果

研究问题

• RQ1 如何在分布式LLM训练中同时防止服务器端和客户端的模型参数与数据泄露？
• RQ2 在TEE保护的分布式LLM框架中，安全、模型准确率与训练效率之间的最优权衡是什么？
• RQ3 一种将后期层隔离在服务器端TEE中的分层微调策略，是否能在减少TEE中处理的参数数量的同时保持高准确率？
• RQ4 与差分隐私或MPC等其他隐私保护技术相比，轻量级加密（OTP）与TEE在性能和安全性方面有何差异？
• RQ5 不同TEE硬件配置（如SGX与TDX/SGX）对训练延迟和系统成本有何影响？

主要发现

• Method2在服务器端TEE中部署4层，QKV和全连接层的参数比例分别为25%和50%，在全部五个医学数据集上均实现了最高平均准确率，优于FL-LLM和SWMT。
• Method2的平均训练时间仅为每数据样本4.33秒（LoRA），显著低于SWMT（17.92秒）和Method1（17.06秒），展现出极高的效率。
• Method1的准确率略低于FL-LLM（明文）是由于加密导致的FP16截断误差，但仍提供了强大的安全保证。
• Method2的可训练参数数量约为Method1的五倍，但性能显著更优，证明了分层微调策略的有效性。
• SWMT由于仅使用TEE的CPU进行处理，训练和推理延迟最高；而Method2通过将大部分计算卸载至GPU，最大限度减少了TEE的使用。
• 消融实验表明，增加TEE中的层数可提升准确率，但会带来更高的延迟和更多参数，验证了4层、(25%,50%)比例的配置为平衡之选。