PUMA: Inferência Segura de LLaMA-7B em Cinco Minutos

PUMA: Secure Inference of LLaMA-7B in Five Minutes

July 24, 2023

Autores: Ye Dong, Wen-jie Lu, Yancheng Zheng, Haoqi Wu, Derun Zhao, Jin Tan, Zhicong Huang, Cheng Hong, Tao Wei, Wenguang Cheng

cs.AI

Resumo

Com o ChatGPT como representante, inúmeras empresas começaram a oferecer serviços baseados em grandes modelos Transformers. No entanto, o uso desses serviços inevitavelmente expõe os prompts dos usuários ao provedor do modelo. Estudos anteriores investigaram a inferência segura para modelos Transformers utilizando computação multipartidária segura (MPC), onde os parâmetros do modelo e os prompts dos clientes são mantidos em sigilo. Apesar disso, essas estruturas ainda são limitadas em termos de desempenho do modelo, eficiência e implantação. Para superar essas limitações, propomos o framework PUMA, que permite uma inferência rápida e segura em modelos Transformers. Nosso framework desenvolve aproximações de alta qualidade para funções computacionalmente caras, como GeLU e Softmax, reduzindo significativamente o custo da inferência segura enquanto mantém o desempenho do modelo. Além disso, projetamos procedimentos seguros para Embedding e LayerNorm que implementam fielmente a funcionalidade desejada sem comprometer a arquitetura Transformer. O PUMA é cerca de 2x mais rápido que o estado da arte em frameworks MPC, como o MPCFORMER (ICLR 2023), e possui precisão semelhante aos modelos em texto plano sem necessidade de ajuste fino (algo que os trabalhos anteriores não conseguiram alcançar). Outro ponto importante é que o PUMA consegue avaliar o LLaMA-7B em aproximadamente 5 minutos para gerar 1 token. Até onde sabemos, esta é a primeira vez que um modelo com esse tamanho de parâmetros pode ser avaliado sob MPC. O PUMA foi disponibilizado como código aberto no repositório Github do SecretFlow-SPU.

English

With ChatGPT as a representative, tons of companies have began to provide services based on large Transformers models. However, using such a service inevitably leak users' prompts to the model provider. Previous studies have studied secure inference for Transformer models using secure multiparty computation (MPC), where model parameters and clients' prompts are kept secret. Despite this, these frameworks are still limited in terms of model performance, efficiency, and deployment. To address these limitations, we propose framework PUMA to enable fast and secure Transformer model inference. Our framework designs high quality approximations for expensive functions, such as GeLU and Softmax, which significantly reduce the cost of secure inference while preserving the model performance. Additionally, we design secure Embedding and LayerNorm procedures that faithfully implement the desired functionality without undermining the Transformer architecture. PUMA is about 2x faster than the state-of-the-art MPC framework MPCFORMER(ICLR 2023) and has similar accuracy as plaintext models without fine-tuning (which the previous works failed to achieve). One more thing, PUMA can evaluate LLaMA-7B in around 5 minutes to generate 1 token. To our best knowledge, this is the first time that a model with such a parameter size is able to be evaluated under MPC. PUMA has been open-sourced in the Github repository of SecretFlow-SPU.