Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Fast, Performant, Secure Distributed Training Framework For Large Language Model

Wei Huang, Yinggui Wang, Anda Cheng, Aihui Zhou, Chaofan Yu, Lei Wang

TL;DR

The paper tackles the problem of parameter and data leakage in distributed (federated) LLM training. It introduces a secure framework based on model slicing that uses TEEs on client and server, plus OTP-based lightweight encryption, with two schemes: Method1 for small-memory TEEs and Method2 with a layer-wise split and SPF to boost accuracy. The main contributions are: (i) a secure distributed training approach combining TEEs and encryption, (ii) a split fine-tuning strategy (SPF) to enhance downstream task performance, and (iii) empirical results on five medical datasets showing strong accuracy under security constraints. This work enables cost-effective, scalable, and secure co-training of domain-specific LLMs on siloed data, offering practical tradeoffs between hardware requirements and performance.

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.

A Fast, Performant, Secure Distributed Training Framework For Large Language Model

TL;DR

The paper tackles the problem of parameter and data leakage in distributed (federated) LLM training. It introduces a secure framework based on model slicing that uses TEEs on client and server, plus OTP-based lightweight encryption, with two schemes: Method1 for small-memory TEEs and Method2 with a layer-wise split and SPF to boost accuracy. The main contributions are: (i) a secure distributed training approach combining TEEs and encryption, (ii) a split fine-tuning strategy (SPF) to enhance downstream task performance, and (iii) empirical results on five medical datasets showing strong accuracy under security constraints. This work enables cost-effective, scalable, and secure co-training of domain-specific LLMs on siloed data, offering practical tradeoffs between hardware requirements and performance.

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.
Paper Structure (12 sections, 3 equations, 2 figures, 4 tables)

This paper contains 12 sections, 3 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. RELATE WORK
  3. Security Threats to Federated Learning
  4. Trusted Execution Environment (TEE)
  5. METHOD
  6. Method1: Small memory (Consumer-grade) TEE-shielded LLM Partition
  7. Method2: Large memory TEE-shielded LLM Partition
  8. EXPERIMENTAL EVALUATION
  9. Datasets and Experimental Configurations
  10. Experimental Results of Method1 and Method2
  11. Ablation Experiment for Method2
  12. CONCLUSION

Figures (2)

  • Figure 1: Block diagram of the distributed LLM based on the model partition for Method1, where the green plus sign indicates encryption and the red minus sign indicates decryption.
  • Figure 2: The distributed LLM training for Method2.