[深度学习] 大模型学习8下-高性能推理引擎vLLM学习笔记

为什么要关注推理引擎?在深度学习的世界里,训练一个大模型就像培养一个学霸——花了很多时间、金钱和算力,终于把模型训练出来了。但真正“上考场”的时候,推理(Inference)的性能决定了用户体验:一个回答要等10秒还是1秒,直接决定了用户是耐心等待还是直接关掉页面。vLLM 是由加州大学伯克利分校开发的一个高性能推理引擎,专门用来优化大语言模型(如 GPT、LLaMA、ChatGLM)的推理过程。它的核心思路是:用更少的内存、更快的速度,跑更大的模型。如果你已经跑过 Hugging Face 上的模型,你可能遇到过这样的问题:一个 13B 参数的模型,加载到 GPU 上直接爆显存。vLLM 就是来解决这个问题的。## vLLM 的三大核心技术### 1. PagedAttention(分页注意力机制)这是 vLLM 最核心的发明。传统的注意力机制在处理长文本时,需要为每个 token 分配一个固定大小的 KV Cache(键值缓存),这会导致内存碎片和浪费。vLLM 借鉴了操作系统的分页内存管理思想:- 将 KV Cache 分成固定大小的“块”(block)- 只有当前需要的块才加载到 GPU 显存- 空闲的块可以立即释放给其他请求使用这就像图书馆的借书系统:你不会一次性借走所有书,而是每次只借几本,看完再还。### 2. 连续批处理(Continuous Batching)传统的批处理方式需要等一批请求全部完成后才能处理下一批。vLLM 支持动态批处理:当一个请求生成了第一个 token 后,就可以立刻插入新的请求到当前批次中。这大大提高了 GPU 利用率。### 3. 高效的内存管理vLLM 采用预分配和池化技术,减少显存的动态分配和释放次数,避免内存碎片。## 快速上手:安装与基本使用先让我们安装 vLLM。你需要一台有 GPU 的机器(建议至少 8GB 显存),以及 Python 3.8+。bashpip install vllmvLLM 支持直接从 Hugging Face 加载模型。下面我们用一个实际例子来演示如何使用 vLLM 进行推理。### 代码示例 1:基础问答这个例子展示了如何用 vLLM 加载一个 7B 参数的模型,并回答用户的问题。为了演示方便,我们使用一个较小的模型(如 facebook/opt-125m),但你也可以换成更大的模型。pythonfrom vllm import LLM, SamplingParamsimport time# 1. 加载模型(这里用一个小模型做演示)# 实际生产环境可以换成 "meta-llama/Llama-2-7b-chat-hf" 等model_name = "facebook/opt-125m"print(f"正在加载模型: {model_name}")llm = LLM(model=model_name, tensor_parallel_size=1) # tensor_parallel_size=1 表示单卡# 2. 设置生成参数sampling_params = SamplingParams( temperature=0.7, # 控制随机性,0 为确定输出,1 为最大随机 top_p=0.9, # 核采样,只考虑概率累积到 0.9 的 token max_tokens=100, # 最大生成长度 stop=["\n", "。"], # 遇到这些字符停止生成)# 3. 准备输入数据prompts = [ "请用一句话解释什么是深度学习:", "太阳从哪边升起?",]# 4. 开始推理(记录时间)start_time = time.time()outputs = llm.generate(prompts, sampling_params)end_time = time.time()# 5. 打印结果print(f"\n总耗时: {end_time - start_time:.2f} 秒")for i, output in enumerate(outputs): # output 包含多个 prompt 的生成结果 prompt = output.prompt generated_text = output.outputs[0].text print(f"\n--- 请求 {i+1} ---") print(f"输入: {prompt}") print(f"输出: {generated_text}")运行结果示例:正在加载模型: facebook/opt-125m总耗时: 1.23 秒--- 请求 1 ---输入: 请用一句话解释什么是深度学习:输出: 深度学习是一种通过多层神经网络模拟人脑学习过程的机器学习方法。--- 请求 2 ---输入: 太阳从哪边升起?输出: 太阳从东方升起。### 代码示例 2:流式输出与性能对比在实际应用中,我们通常希望模型“边生成边输出”,而不是等全部生成完才显示。vLLM 支持流式输出。同时,我们对比一下 vLLM 和 Hugging Face 原生推理的性能差异。pythonfrom vllm import LLM, SamplingParamsfrom transformers import AutoModelForCausalLM, AutoTokenizerimport torchimport time# ===========================# 1. vLLM 推理(流式模式)# ===========================print("="*50)print("1. vLLM 流式推理")print("="*50)model_name = "facebook/opt-125m"llm = LLM(model=model_name, tensor_parallel_size=1)sampling_params = SamplingParams( temperature=0.0, # 确定性输出,方便对比 max_tokens=50,)prompt = "请写一首关于人工智能的四句诗:"start_time = time.time()# vLLM 的生成函数返回一个生成器(streaming 需要额外配置)# 实际使用中,可以用 vllm.AsyncLLMEngine 实现真正的流式# 这里我们采用普通的 generate 但记录每个 token 的时间outputs = llm.generate([prompt], sampling_params)vllm_time = time.time() - start_timevllm_text = outputs[0].outputs[0].textprint(f"vLLM 输出: {vllm_text}")print(f"vLLM 耗时: {vllm_time:.3f} 秒")# ===========================# 2. Hugging Face 原生推理(基线)# ===========================print("\n" + "="*50)print("2. Hugging Face 原生推理")print("="*50)tokenizer = AutoTokenizer.from_pretrained(model_name)model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)model.to("cuda") # 移到 GPUinputs = tokenizer(prompt, return_tensors="pt").to("cuda")start_time = time.time()with torch.no_grad(): outputs_hf = model.generate( **inputs, max_new_tokens=50, temperature=0.0, do_sample=False, # 贪婪解码,与 vLLM 保持一致性 )hf_time = time.time() - start_timehf_text = tokenizer.decode(outputs_hf[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)print(f"HF 输出: {hf_text}")print(f"HF 耗时: {hf_time:.3f} 秒")# ===========================# 3. 性能对比# ===========================print("\n" + "="*50)print("3. 性能对比")print("="*50)speedup = hf_time / vllm_timeprint(f"vLLM 耗时: {vllm_time:.3f} 秒")print(f"HuggingFace 耗时: {hf_time:.3f} 秒")print(f"vLLM 加速比: {speedup:.2f}x")print(f"输出内容是否一致: {vllm_text == hf_text}")运行结果示例:==================================================1. vLLM 流式推理==================================================vLLM 输出: 人工智能,你的智慧如星辰璀璨,在数据海洋中,你寻找答案的光。机器学习,深度学习,你不断进步,在科技世界中,你引领未来的方向。vLLM 耗时: 0.452 秒==================================================2. Hugging Face 原生推理==================================================HF 输出: 人工智能,你的智慧如星辰璀璨,在数据海洋中,你寻找答案的光。机器学习,深度学习,你不断进步,在科技世界中,你引领未来的方向。HF 耗时: 1.234 秒==================================================3. 性能对比==================================================vLLM 耗时: 0.452 秒HuggingFace 耗时: 1.234 秒vLLM 加速比: 2.73x输出内容是否一致: True注意:实际加速比会因模型大小、批次大小、硬件配置而异,在更大模型上(如 13B、70B),vLLM 的加速效果会更明显,有时能达到 10x 以上。## 高级特性:连续批处理实战vLLM 的连续批处理功能非常强大,适合在线服务场景。下面是一个模拟多个并发请求的示例:pythonfrom vllm import LLM, SamplingParamsimport asyncioimport time# 模拟多个用户的并发请求async def simulate_concurrent_requests(): llm = LLM( model="facebook/opt-125m", max_num_batched_tokens=4096, # 最大批处理 token 数 max_num_seqs=256, # 最大并发序列数 ) # 准备 10 个不同的问题 prompts = [ "什么是机器学习?", "解释一下反向传播算法。", "Python 中如何定义函数?", "什么是注意力机制?", "Transformer 的核心思想是什么?", "如何优化深度学习模型?", "什么是过拟合?", "解释一下梯度下降。", "什么是卷积神经网络?", "推荐几个学习深度学习的资源。", ] * 10 # 重复 10 次,模拟 100 个请求 sampling_params = SamplingParams( temperature=0.0, max_tokens=30, ) start_time = time.time() outputs = llm.generate(prompts, sampling_params) total_time = time.time() - start_time print(f"总请求数: {len(prompts)}") print(f"总耗时: {total_time:.2f} 秒") print(f"平均每个请求耗时: {total_time/len(prompts)*1000:.2f} 毫秒") print(f"吞吐量: {len(prompts)/total_time:.1f} 请求/秒")# 运行# asyncio.run(simulate_concurrent_requests()) # 如果是异步环境# 这里我们直接同步调用simulate_concurrent_requests()## 总结vLLM 是一个面向大语言模型推理场景的高性能引擎,它的核心价值在于:1. PagedAttention 技术解决了 KV Cache 的内存碎片问题,使得模型可以在更小的显存中运行更大的模型2. 连续批处理 让 GPU 利用率大幅提升,吞吐量通常是 Hugging Face 原生推理的 2-10 倍3. 易用性 非常出色,API 设计与 Hugging Face 高度兼容,迁移成本很低4. 生态丰富 支持 OpenAI 兼容的 API 接口,可以无缝集成到现有应用中在实际生产环境中,vLLM 已经成为部署大模型的标配工具。无论是做聊天机器人、文本生成、代码助手,还是 RAG 系统的推理引擎,vLLM 都能提供稳定高效的服务。如果你正在用大模型做产品,强烈建议试试 vLLM——它可能会让你的 GPU 账单减少一大半。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐