不可思议！AirLLM 如何让 70B 大模型在 4GB GPU 上顺利推理？

发布日期：2024-11-22 06:02:39 浏览次数： 4112

作者：小兵的AI视界

微信搜一搜，关注“小兵的AI视界”

在人工智能蓬勃发展的浪潮中，大型语言模型（LLM）无疑是最耀眼的明珠。它们在自然语言处理、内容生成、智能问答等众多领域展现出了惊人的能力。

然而，这些强大的模型在推理阶段往往对计算资源和内存有着极高的要求，这成为了其广泛应用的一大瓶颈。许多研究人员和开发者受限于硬件条件，无法充分挖掘大模型的潜力。

在此背景下，AirLLM应运而生，它是一款专门针对大型语言模型推理阶段内存使用优化的开源工具包，为在资源受限环境下实现大模型的高效推理开辟了新的途径一、技术亮点

1、分层推理

AirLLM的分层推理基于分而治之的理念。大型语言模型通常由一个嵌入投影层和多个相同的Transformer层构成。在推理过程中，各层顺序执行，上一层输出作为下一层输入。AirLLM利用这一特性，在执行某一层推理时，仅从磁盘加载该层参数，计算完成后释放内存。例如，对于70B参数的大模型，每层只需约1.6GB GPU内存（约为整个模型的1/80）。即便考虑输出缓存（如KV缓存），当输入长度为100时，额外GPU内存仅约30MB，这使得整个70B模型推理在4GB显存内即可完成。

2、FlashAttention

FlashAttention是AirLLM性能提升的关键技术。它源于对self-attention内存复杂度优化的研究，将原始O(n²)的内存复杂度降低至O(n)。通过按顺序计算和更新中间结果，并丢弃不必要数据的方式，实现了这一突破。同时，FlashAttention深度优化了CUDA内存访问，从而在推理和训练阶段实现多倍加速。这种优化使得模型在处理长序列文本时更加高效，显著提升了推理速度。

3、模型文件共享

原始大型语言模型文件通常分割为较大块（如10GB左右），但分层推理每层仅需约1.6GB数据。若按原始块加载，每层执行都要重新加载整个文件，会造成大量内存浪费。AirLLM通过预处理和分层分割原始Hugging Face模型文件，并使用安全张量（SafeTensors）存储，以内存映射方式加载。SafeTensors确保存储和内存格式匹配，加快磁盘读取速度，减少重复加载和磁盘读取导致的内存损耗，有效提高了推理效率。

4、元设备

AirLLM借助Hugging Face Accelerate的元设备（MetaDevice）功能优化内存。元设备是为运行超大型模型设计的虚拟设备，通过元设备加载模型时，仅加载代码，内存使用率为零。在执行过程中，根据需要将模型部分内容从元设备动态转移到CPU或GPU等真实设备。这种动态加载机制使模型能根据硬件资源灵活调整，在保证推理性能的同时，最大程度减少内存占用。

二、功能特点

1、低内存需求

AirLLM的突出优势在于其低内存需求。它打破了大型语言模型对高端硬件的依赖，能让70亿参数级别的大型语言模型在仅4GB VRAM的单张GPU卡上运行，且无需量化、蒸馏或剪枝等模型压缩手段。更厉害的是，它可使405B参数的Llama 3.1模型在8GB VRAM的硬件上运行，极大降低了大模型推理的硬件门槛，让更多设备能运行大模型。

2、高度灵活性

AirLLM适用于多种场景，无论是学术研究中的模型探索、工业生产中的应用开发，还是在移动设备、资源受限的数据中心部署AI应用，都能表现出色。它能根据不同硬件条件和应用场景自适应调整模型运行方式，满足多样化需求，推动AI在各领域广泛应用。

3、开源特性

作为开源项目，AirLLM吸引了全球开发者参与。这种开源模式促进了技术交流与协作，开发者可贡献代码、提出改进建议，共同推动项目发展。同时，开源也便于开发者根据自身需求定制和扩展功能，加速了大型语言模型推理技术的创新。

三、应用场景

1、教育科研领域

在教育和科研中，AirLLM意义重大。过去，探索先进NLP模型需要昂贵硬件，限制了高校实验室和个人研究者。现在，借助AirLLM，他们在本地普通设备上就能运行70B甚至更大参数的模型，降低科研门槛，激发科研创新，推动自然语言处理研究发展，培养更多AI人才。

2、企业级应用开发

对于企业开发团队，AirLLM助力快速实现AI功能。在构建智能客服、内容生成平台、数据分析预测等应用时，可基于现有硬件资源利用AirLLM运行大型语言模型，加速企业数字化转型和AI应用落地，提升企业竞争力和创新能力。

3、个人开发者创新

个人开发者也因AirLLM受益。他们能在笔记本电脑等普通硬件上运行大规模模型，进行测试、开发和创新。这为个人开发者提供广阔创作空间，使其能自由探索AI应用边界，开发如写作助手、智能学习工具等创新应用。

四、使用方法

1、安装 airllm 包

通过`pip install airllm`命令安装AirLLM包。

pip install airllm

2、模型推理示例（以Llama 2模型为例）

from airllm import AutoModel
定义模型输入的最大长度MAX_LENGTH = 128
使用预训练模型的路径或名称来加载模型这里使用的是Hugging Face模型库中的模型IDmodel = AutoModel.from_pretrained("garage-bAInd/Platypus2-70B-instruct")
如果模型已经下载到本地，也可以使用本地路径来加载模型这一行被注释掉了，如果需要使用本地模型，可以取消注释#model = AutoModel.from_pretrained("/home/ubuntu/.cache/huggingface/hub/models--garage-bAInd--Platypus2-70B-instruct/snapshots/b585e74bcaae02e52665d9ac6d23f4d0dbc81a0f")
定义输入文本列表，这里只有一个问题input_text = ['What is the capital of United States?',#'I like',]
使用模型的tokenizer将输入文本转换为模型可以理解的tokenreturn_tensors="pt" 表示返回PyTorch张量return_attention_mask=False 表示不返回注意力掩码truncation=True 表示如果输入超过最大长度则截断max_length=MAX_LENGTH 设置最大长度限制padding=False 表示不进行填充input_tokens = model.tokenizer(input_text,return_tensors="pt", return_attention_mask=False, truncation=True, max_length=MAX_LENGTH, padding=False)
使用模型生成文本input_tokens['input_ids'].cuda() 将输入ID传递到GPUmax_new_tokens=20 设置生成的最大新token数量use_cache=True 表示使用缓存来加速生成过程return_dict_in_generate=True 表示返回一个字典而不是直接返回生成的序列generation_output = model.generate(input_tokens['input_ids'].cuda(), max_new_tokens=20,use_cache=True,return_dict_in_generate=True)
将生成的token序列解码回文本output = model.tokenizer.decode(generation_output.sequences[0])
打印生成的文本print(output)

‍五、结语

AirLLM作为大型语言模型推理领域的创新开源工具包，凭借其卓越的技术亮点、突出的功能特点和广泛的应用场景，为解决大模型在资源受限环境下的推理难题提供了有效方案。它推动了AI技术的普及和发展，让更多人能够利用大型语言模型创造价值。无论是科研人员、企业开发者还是个人开发者，都能从AirLLM中获得助力，开启更多的可能性。