LLM的轻量化推理：AirLLM

AirLLM是一种针对LLM的轻量级推理框架，通常用于优化和加速大模型的推理过程，可以允许70B的LLM在单个4GB的GPU上运行，无需量化、蒸馏、剪枝。

AirLLM是通过分层推理实现的上述功能，本质上是将LLM的每一层拆分出来，分别加载权重进行推理，在推理较大的LLM时会很慢。

环境配置

使用pip安装airllm包

pip install airllm

环境需要的其他package与正常推理部署LLM的一致

模型推理

在模型的第一次加载过程中会将模型权重逐层拆分保存，大小与原权重大小相同，请确保有足够的磁盘空间

from airllm import AutoModel

MAX_LENGTH = 128

model = AutoModel.from_pretrained(model_name/model_dir)

input_text = [
        "who is the president of US?"
]

input_tokens = model.tokenizer(
        input_text,
        return_tensors="pt",
        return_attention_mask=False,
        truncation=True,
        max_length=MAX_LENGTH,
        padding=False
)    

generation_output = model.generate(
        input_tokens["input_ids"].cuda(),
        max_new_tokens=20,
        use_cache=True,
        return_dict_in_generate=True
)

output = model.tokenizer.decode(generation_output.sequences[0])

print(output)