[論文レビュー] 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

ひとこと要約

本論文は、モデルスライシング、信頼実行環境（TEE）、および軽量暗号化を用いて、サーバー側およびクライアント側のモデル／データ漏洩を防止する、安全で高性能な大規模言語モデル（LLM）の分散学習フレームワークを提案する。クライアントおよびサーバーの両方に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内で保存および計算され、ファインチューニング済み重みおよび埋め込みの機密性が保証される。
• TEEと外部環境間の安全な通信は、勾配および中間出力にワンタイムパッド（OTP）暗号化を用いることで実現される。
• スプリットファインチューニング戦略が導入され、サーバー側のTEEで更新されるのは層のサブセット（例：最後の4層）に限定される。これにより、パラメータ更新の表面積とコストが低減される。
• 本手法は2種類のTEE構成をサポートする：Intel SGX（小メモリ、GPU非対応）およびIntel TDX/SGX（大メモリ、GPU非対応）、それぞれに最適化されたスキームが用意される。
• スパarsification Parameter Fine-tuning（SPF）をLoRAと組み合わせることで、モデルサイズを増加させずに下流タスクの精度をさらに向上できる。

実験結果

リサーチクエスチョン

• RQ1 分散LLM学習において、サーバー側およびクライアント側の両方のモデルパラメータおよびデータ漏洩をどのように防止できるか？
• RQ2 TEE保護された分散LLMフレームワークにおいて、セキュリティ、モデル精度、学習効率の最適なトレードオフは何か？
• RQ3 後段の層をサーバー側のTEEに隔離するスプリットファインチューニング戦略は、パラメータ処理量を削減しながらも高い精度を維持できるか？
• RQ4 軽量暗号化（OTP）およびTEEの使用は、微分プライバシーまたはMPCなどの他のプライバシー保護技術と比較して、パフォーマンスおよびセキュリティ面でどのように異なるか？
• RQ5 異なるTEEハードウェア構成（例：SGX対TDX/SGX）は、学習遅延およびシステムコストにどのような影響を及えるか？

主な発見

• Method2は、QKVおよびドライヴ層にそれぞれ25%および50%のパラメータ比を適用し、サーバー側のTEEに4層を配置したが、5つの医療データセットすべてで最高の平均精度を達成し、FL-LLMおよびSWMTを上回った。
• Method2の平均学習時間は1データサンプルあたり4.33秒（LoRA）であり、SWMT（17.92秒）およびMethod1（17.06秒）よりも顕著に短く、高い効率性を示した。
• Method1は、暗号化によるFP16丸め誤差の影響でFL-LLM（プレーンテキスト）よりわずかに精度が低かったが、依然として強固なセキュリティ保証を提供した。
• Method2のトレーナブルパラメータ数はMethod1の約5倍であったが、顕著な性能向上を達成しており、スプリットファインチューニングの有効性が裏付けられた。
• SWMTは、TEE専用CPU処理のため、最も高い学習および推論遅延を示した。一方、Method2はGPUに大部分の計算をオフロードすることで、TEEの使用を最小限に抑え、効率を最大化した。
• アブレーションスタディにより、TEE内に配置する層数を増やすと精度は向上するが、遅延およびパラメータ数の増加を伴うことが確認され、4層、(25%,50%)比がバランスの取れた選択であることが妥当性を確認した。