LayerSkip – Meta推出加速大型语言模型推理过程的技术

我们提出的 LayerSkip 是一种端到端的解决方案,可加快大型语言模型(LLM)的推理速度。 首先,在训练过程中,我们采用了层间丢弃技术(layer dropout),早期层间丢弃率较低,后期层间丢弃率较高。 其次,在推理过程中,我们证明这种训练方法提高了早期退出的准确性,而无需在模型中添加任何辅助层或模块。 第三,我们提出了一种新颖的自推测解码方案,即在早期层退出,并通过模型的其余层进行验证和校正。 与其他推测式解码方法相比,我们提出的自推测式解码方法占用的内存更少,并能从草稿和验证阶段的共享计算和激活中获益。 我们在不同大小的 Llama 模型上进行了不同类型的训练实验:从头开始预训练、持续预训练、在特定数据域上进行微调,以及在特定任务上进行微调。 我们实施了推理解决方案,结果表明,CNN/DM 文档的摘要速度提高了 2.16 倍,编码速度提高了 1.82 倍,TOPv2 语义解析任务的速度提高了 2.0 倍。 我们在 https://github.com/facebookresearch/LayerSkip 开源了我们的代码。

快速上手

bash 复制代码
$ git clone git@github.com:facebookresearch/LayerSkip.git
$ cd LayerSkip

创建环境

bash 复制代码
$ conda create --name layer_skip python=3.10
$ conda activate layer_skip

$ pip install -r requirements.txt

访问模型: 为了观察加速情况,您需要访问使用 LayerSkip 配方训练过的 LLM。 我们在 HuggingFace 上提供了 6 个检查点,它们是使用 LayerSkip 配方持续预训练的不同 Llama 模型:

代码

python 复制代码
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from copy import deepcopy

checkpoint = "facebook/layerskip-llama3.2-1B"
early_exit = 4
device = "cuda" if torch.cuda.is_available() else "cpu"
prompt = "typing import List\ndef bucket_sort(A: List):"

model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", use_safetensors=True, torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)

generation_config = model.generation_config

weights_memo = {id(w): w for w in model.parameters()}
assistant_model = deepcopy(model, memo=weights_memo) # Clone main model with shared weights
assistant_model.model.layers = assistant_model.model.layers[:early_exit] # Apply early exit
del assistant_model.model.layers[early_exit:]

inputs = tokenizer(prompt, return_tensors="pt").to(device)

outputs = model.generate(**inputs, generation_config=generation_config, assistant_model=assistant_model, max_new_tokens=512)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True)[0])

或者Torchrun

bash 复制代码
$ torchrun generate.py --model facebook/layerskip-llama2-7B \
    --sample True \
    --max_steps 512

LayerSkip的项目地址

感谢大家花时间阅读我的文章,你们的支持是我不断前进的动力。期望未来能为大家带来更多有价值的内容,请多多关注我的动态!

相关推荐
音视频牛哥43 分钟前
打通视频到AI的第一公里:轻量RTSP服务如何重塑边缘感知入口?
人工智能·计算机视觉·音视频·大牛直播sdk·机器视觉·轻量级rtsp服务·ai人工智能
Wendy14412 小时前
【灰度实验】——图像预处理(OpenCV)
人工智能·opencv·计算机视觉
中杯可乐多加冰2 小时前
五大低代码平台横向深度测评:smardaten 2.0领衔AI原型设计
人工智能
无线图像传输研究探索2 小时前
单兵图传终端:移动场景中的 “实时感知神经”
网络·人工智能·5g·无线图传·5g单兵图传
Aronup2 小时前
NLP学习开始01-线性回归
学习·自然语言处理·线性回归
zzywxc7873 小时前
AI在编程、测试、数据分析等领域的前沿应用(技术报告)
人工智能·深度学习·机器学习·数据挖掘·数据分析·自动化·ai编程
铭keny4 小时前
YOLOv8 基于RTSP流目标检测
人工智能·yolo·目标检测
墨尘游子4 小时前
11-大语言模型—Transformer 盖楼,BERT 装修,RoBERTa 直接 “拎包入住”|预训练白话指南
人工智能·语言模型·自然语言处理
金井PRATHAMA4 小时前
主要分布于内侧内嗅皮层的层Ⅲ的网格-速度联合细胞(Grid × Speed Conjunctive Cells)对NLP中的深层语义分析的积极影响和启示
人工智能·深度学习·神经网络·机器学习·语言模型·自然语言处理·知识图谱
天道哥哥4 小时前
InsightFace(RetinaFace + ArcFace)人脸识别项目(预训练模型,鲁棒性很好)
人工智能·目标检测