完整训练代码 + vLLM 生产级部署方案,32G 消费级显卡全流程跑通!


想拥有自己的专属 AI 模型,却被动辄数百 GB 的显存需求劝退?

今天这篇教程,带你用 LoRA 微调 + vLLM 推理部署,在 32G 消费级显卡 上完成从训练到上线的全流程。

文章包含 完整代码逐行详细解释,以及 生产级部署方案,收藏这一篇就够了。


整体流程概览

加载数据集 → 数据预处理 → 加载基座模型 → 配置LoRA → 训练 → 保存 → 清理显存 → 合并模型 → vLLM部署上线

我们一步步来。


一、环境准备

pip install transformers datasets peft accelerate torch vllm

核心依赖说明:

  • transformers:Hugging Face 模型框架,训练主力
  • datasets:数据处理工具
  • peft:参数高效微调库(LoRA 就在这里)
  • accelerate:分布式训练与显存优化
  • torch:PyTorch 深度学习框架
  • vllm:高性能推理引擎,部署用

二、加载数据集

from datasets import load_datasetimport transformersds = load_dataset('json', data_files='alpaca_gpt4_data_zh.json')ds = ds['train']print(ds[:3])print(transformers.__version__)

详细解释

load_dataset('json', data_files='alpaca_gpt4_data_zh.json')

从 JSON 文件加载中文指令微调数据集。每条数据包含三个字段:

字段 含义 示例
instruction 指令 “请解释什么是光合作用”
input 输入(可选) 提供任务的上下文或具体输入
output 期望输出 “光合作用是植物利用光能…”

ds['train'] 取训练集。

print(ds[:3]) 打印前 3 条数据,快速验证格式。

print(transformers.__version__) 记录版本号,方便日后复现。


三、数据预处理(核心步骤)

3.1 加载分词器

from transformers import AutoTokenizertokenizer = AutoTokenizer.from_pretrained('/home/will/models/llama')tokenizer.pad_token = tokenizer.eos_token
  • AutoTokenizer.from_pretrained(...):加载与 LLaMA 模型配套的分词器,把文本转成模型能理解的 token IDs
  • tokenizer.pad_token = tokenizer.eos_token:LLaMA 原始分词器没有定义 pad_token,借用 eos_token(结束标记)来补齐 batch 内的短序列

3.2 构造训练样本

def process_func(example):    MAX_LENGTH = 512        instruction = tokenizer(        "\n".join(["Human: " + example["instruction"], example["input"]]).strip() + "\n\nAssistant: ",        add_special_tokens=False    )    response = tokenizer(        example["output"] + tokenizer.eos_token,        add_special_tokens=False    )        input_ids = instruction["input_ids"] + response["input_ids"]    attention_mask = instruction["attention_mask"] + response["attention_mask"]    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"]        if len(input_ids) > MAX_LENGTH:        input_ids = input_ids[:MAX_LENGTH]        attention_mask = attention_mask[:MAX_LENGTH]        labels = labels[:MAX_LENGTH]        return {        "input_ids": input_ids,        "attention_mask": attention_mask,        "labels": labels    }
这段代码是整篇文章最需要理解的部分

📌 对话格式构造

把数据拼接成如下格式:

Human: 请解释什么是光合作用Assistant: 光合作用是植物利用光能...<eos>

Human: 代表用户提问,Assistant: 代表模型回答。这种格式让模型学会"问答对话"模式。

📌 add_special_tokens=False

默认分词器会在文本开头自动加 <bos>。设为 False 是因为我们要手动控制 token 拼接顺序——先 instruction 后 response,由自己决定特殊标记的位置。

📌 labels 的设置(指令微调的精髓!)

labels = [-100] * len(instruction["input_ids"]) + response["input_ids"]
  • -100 是 CrossEntropyLoss 的 ignore_index,标记为 -100 的位置 不计算损失、不参与训练
  • 只有 Assistant: 之后的回复部分才参与训练

为什么? 因为我们要训练的是模型 “如何回答”,而不是训练它 “如何重复问题”。

完整 input_ids:  [Human: ... \n\nAssistant: ...回复内容...]                 |←────── 忽略不训练 ──────→|←── 训练这部分 ──→|对应 labels:     [-100, -100, ..., -100,  回复的token_ids...]

📌 MAX_LENGTH = 512

限制每条样本最大 512 个 token。超长截断,目的是控制显存——序列越长,显存消耗越大。32G 显卡处理 512 长度比较安全。

3.3 执行预处理

tokenized_ds = ds.map(process_func, remove_columns=ds.column_names)print(tokenized_ds)print(tokenizer.decode(tokenized_ds[0]["input_ids"]))print(tokenizer.decode(list(filter(lambda x: x != -100, tokenized_ds[0]["labels"]))))
  • ds.map(...):对每条样本应用处理函数
  • remove_columns=...:删除原始列,只保留 input_idsattention_masklabels
  • 下面三行 验证处理结果:打印数据结构、还原 input_ids 为文本、过滤 -100 后还原 labels,确保格式正确

四、加载基座模型

import torchfrom transformers import AutoModelForCausalLMmodel = AutoModelForCausalLM.from_pretrained(    "/home/will/models/llama",    low_cpu_mem_usage=True,    torch_dtype=torch.float16,    device_map="auto",    use_cache=False)

参数详解

参数 为什么这么设
low_cpu_mem_usage True 减少 CPU 内存占用,避免加载大模型时 OOM
torch_dtype float16 半精度,显存比 float32 减半,32G 显卡必备
device_map "auto" 自动将模型层分配到可用 GPU
use_cache False 训练时不需要推理 KV cache,关掉省显存

💡 8bit 量化选项(代码注释中提供)

# model = AutoModelForCausalLM.from_pretrained(#     "/home/will/models/llama",#     load_in_8bit=True,#     device_map="auto",#     use_cache=False# )

如果模型较大(如 13B),建议开启 8bit 量化,需安装 bitsandbytes


五、配置 LoRA 微调

from peft import LoraConfig, TaskType, get_peft_modelconfig = LoraConfig(task_type=TaskType.CAUSAL_LM)model = get_peft_model(model, config)model.print_trainable_parameters()

LoRA 原理一句话

不改动原始权重,在每层旁边加一个小型"旁路适配器",只训练这个小适配器。

就像改造大楼不需要重建,只需要在每层加个小阳台——阳台很小,成本自然低。

LoraConfig 默认参数

参数 默认值 含义
r 8 秩(rank),决定适配器大小。越大表达能力越强
lora_alpha 8 缩放系数,控制 LoRA 权重影响力
lora_dropout 0 Dropout 比率,防过拟合
target_modules None 自动检测所有线性层(q_proj, v_proj 等)

生产环境建议显式指定:

config = LoraConfig(    task_type=TaskType.CAUSAL_LM,    r=16,    lora_alpha=32,       # 通常是 r 的 2 倍    lora_dropout=0.05,    target_modules=["q_proj", "v_proj"],)

get_peft_model 做了什么?

  • 原始模型权重被 冻结(不参与训练,不需要梯度)
  • 只有 LoRA 适配器的参数 可训练
  • 可训练参数通常只有原模型的 0.1% ~ 1%

执行 print_trainable_parameters() 你会看到类似:

trainable params: 4,194,304 || all params: 7,000,000,000 || trainable%: 0.0599%

7B 的模型,只训练约 400 万参数!


六、配置训练参数

from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seqargs = TrainingArguments(    output_dir="./chatbot",    per_device_train_batch_size=1,    per_device_eval_batch_size=1,    gradient_accumulation_steps=8,    logging_steps=10,    num_train_epochs=1,    gradient_checkpointing=True,    optim="adamw_torch",    save_steps=500,    save_total_limit=2,    fp16=True,    bf16=False,    dataloader_num_workers=0,    remove_unused_columns=False,)

参数逐行详解

参数 含义
output_dir "./chatbot" 模型和检查点保存目录
per_device_train_batch_size 1 每张 GPU 每次处理 1 条,32G 显卡安全值
gradient_accumulation_steps 8 梯度累积 。等效 batch_size = 1 × 8 = 8。每 8 步才更新一次权重
logging_steps 10 每 10 步打印训练日志
num_train_epochs 1 指令微调 1-3 轮即可,多了容易过拟合
gradient_checkpointing True 梯度检查点 。用时间换显存,省约 60% 激活显存
optim "adamw_torch" AdamW 优化器
save_steps 500 每 500 步保存检查点
save_total_limit 2 最多保留 2 个检查点,超了自动删最旧的
fp16 True 混合精度训练
bf16 False 不启用 bf16(需 Ampere 架构 RTX 30/40 系列)
dataloader_num_workers 0 主进程加载,避免多进程问题
remove_unused_columns False 保留所有数据列

关于有效 Batch Size

有效 batch_size = per_device_batch_size × gradient_accumulation_steps × GPU数量               = 1 × 8 × 1 = 8

虽然每次只处理 1 条数据,但通过梯度累积,等效于 batch_size=8 的训练效果。


七、开始训练

trainer = Trainer(    model=model,    args=args,    train_dataset=tokenized_ds.select(range(6000)),    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True))import gcgc.collect()torch.cuda.empty_cache()trainer.train()

关键点

  • tokenized_ds.select(range(6000)):取前 6000 条样本。大数据集可先取一部分跑通流程
  • DataCollatorForSeq2Seq:batch 内序列对齐,自动 padding
  • gc.collect() + torch.cuda.empty_cache():训练前清理显存碎片
  • trainer.train():一行启动,自动处理前向传播、反向传播、梯度累积、学习率调度、日志记录、检查点保存

训练日志示例:

{'loss': 1.2345, 'learning_rate': 5e-5, 'epoch': 0.5}{'loss': 1.1234, 'learning_rate': 4.8e-5, 'epoch': 0.6}

loss 逐渐下降,模型在认真学习!


八、保存模型

trainer.save_model()tokenizer.save_pretrained("./chatbot")

这里保存的 只是 LoRA 适配器的权重(通常只有几十 MB),不是完整模型。

好处:同一基础模型可以训练多个不同任务的适配器,灵活切换。


九、清理显存(为合并做准备)

del modeldel trainertorch.cuda.empty_cache()gc.collect()

删除训练对象,释放显存。下一步的模型合并需要额外空间,不清理会 OOM。


十、合并模型

base_model = AutoModelForCausalLM.from_pretrained(    "/home/will/models/llama",    torch_dtype=torch.float16,    device_map="cpu",    low_cpu_mem_usage=True)from peft import PeftModelmodel = PeftModel.from_pretrained(base_model, "./chatbot")merged_model = model.merge_and_unload()merged_output_dir = "./chatbot_merged"merged_model.save_pretrained(merged_output_dir)tokenizer.save_pretrained(merged_output_dir)print(f"Merged model saved to {merged_output_dir}")print("Training and merge completed!")

为什么要合并?

LoRA 训练后权重是 “分离” 的——基础模型一套,适配器一套。合并把 LoRA 权重加回基础模型对应层,得到完整的独立模型。

为什么加载到 CPU?

device_map="cpu"
  • 合并不需要 GPU 加速
  • 把 GPU 腾出来避免显存不足
  • 32G 显卡同时加载基础模型 + 做合并操作压力太大

合并后的目录结构

./chatbot_merged/├── config.json           # 模型配置├── pytorch_model.bin     # 完整模型权重├── tokenizer.json        # 分词器文件├── tokenizer_config.json└── ...

十一、vLLM 部署上线

模型训练合并完成后,就可以部署了。这里我们使用 vLLM——目前最快的开源大模型推理引擎。

启动命令

python -m vllm.entrypoints.openai.api_server \    --model /code/chatbot_merged \    --served-model-name llama \    --max-model-len 8192 \    --host 0.0.0.0 \    --port 6006 \    --dtype bfloat16 \    --gpu-memory-utilization 0.8 \    --trust-request-chat-template \    --enable-auto-tool-choice \    --tool-call-parser hermes

参数逐行详解

参数 含义
--model /code/chatbot_merged 指向合并后的模型目录
--served-model-name llama API 中显示的模型名称
--max-model-len 8192 最大上下文长度 8K 。模型能处理的最长 token 数
--host 0.0.0.0 监听所有网络接口,允许远程访问
--port 6006 服务端口
--dtype bfloat16 使用 bf16 精度推理 。比 fp16 数值稳定性更好,不会溢出(前提是你的显卡支持 bf16,RTX 30/40 系列都支持)
--gpu-memory-utilization 0.8 限制 GPU 显存使用率为 80% 。留 20% 给系统和其他进程,避免把显卡吃满导致崩溃
--trust-request-chat-template 信任模型自带的对话模板,自动处理 Human/Assistant 格式
--enable-auto-tool-choice 启用自动工具选择,让模型能够自主决定何时调用外部工具
--tool-call-parser hermes 使用 Hermes 格式解析工具调用。如果你的模型支持 Function Calling,vLLM 会自动识别并执行

为什么用 vLLM?

相比直接用 transformers 推理,vLLM 有巨大优势:

特性 transformers vLLM
推理速度 基准 快 2-10 倍
显存效率 一般 PagedAttention 技术,显存利用率极高
并发处理 支持高并发,自带批处理
API 兼容性 需要自己写 兼容 OpenAI API 格式

最关键的一点:vLLM 的 API 完全兼容 OpenAI 格式。 这意味着所有用 OpenAI SDK 的代码,只需要改一行就能切换到自己的模型:

# 只需要改 base_url,其他代码不用动!from openai import OpenAIclient = OpenAI(    base_url="http://localhost:6006/v1",    api_key="not-needed"  # 本地部署不需要 key)response = client.chat.completions.create(    model="llama",    messages=[        {"role": "user", "content": "请介绍一下中国的首都"}    ])print(response.choices[0].message.content)

验证服务是否正常

启动后访问 http://你的IP:6006/v1/models,如果返回模型信息说明服务正常。

也可以用 curl 测试:

curl http://localhost:6006/v1/chat/completions \  -H "Content-Type: application/json" \  -d '{    "model": "llama",    "messages": [      {"role": "user", "content": "你好,请自我介绍一下"}    ]  }'

关于 dtype 的补充说明

训练时用的是 fp16,部署用的是 bf16。这是可以的:

  • 模型权重在合并后保存在磁盘上,加载时 vLLM 会自动转换
  • bf16fp16 数值范围更大,推理时更稳定,不容易出现 NaN
  • 前提条件:你的显卡需要支持 bf16(RTX 3090/4090 等 Ampere/Ada 架构都支持)

如果你的显卡不支持 bf16(如 RTX 2080Ti),改成 --dtype float16 即可。


十二、32G 显卡优化总结

训练阶段

技巧 效果
float16 半精度 显存减半
gradient_checkpointing 省约 60% 激活显存
batch_size=1 + gradient_accumulation=8 小 batch 模拟大 batch
MAX_LENGTH=512 控制序列长度
训练后 del + empty_cache 及时释放显存
模型合并放在 CPU 避免 GPU OOM

推理阶段

技巧 效果
--gpu-memory-utilization 0.8 限制显存使用,留出余量
--dtype bfloat16 更稳定的半精度推理
--max-model-len 8192 支持长上下文(训练时 512 也没关系,推理时可以更长)
vLLM PagedAttention 显存利用率远高于原生推理

十三、常见问题

Q1: 训练很慢怎么办?

  • 先用少量数据(1000 条)跑通流程
  • 降低 MAX_LENGTH 到 256
  • 确认模型确实跑在 GPU 上

Q2: loss 不下降?

  • 添加 learning_rate=2e-4 到 TrainingArguments
  • tokenizer.decode 检查数据格式
  • 尝试 num_train_epochs=3

Q3: 合并时 OOM?

  • 确保训练后已 del model 并清理显存
  • device_map="cpu" 必须设

Q4: vLLM 启动报错 Out of Memory?

  • 降低 --gpu-memory-utilization 到 0.7 或 0.6
  • 减小 --max-model-len
  • 确认没有其他进程占用 GPU(nvidia-smi 查看)

Q5: 训练用 fp16,部署用 bf16 会不会有问题?

  • 不会。合并后的模型权重加载时 vLLM 会自动转换精度
  • bf16 比 fp16 数值范围更大,推理更稳定

十四、总结

回顾整个流程:

  1. 加载数据集 → JSON 格式指令数据
  2. 数据预处理 → 构造对话格式,labels 只对回复部分计算损失
  3. 加载基座模型 → 半精度 + 自动设备分配
  4. 配置 LoRA → 只训练 0.1% 参数
  5. 训练 → Trainer 一行启动
  6. 保存 + 合并 → 得到完整模型
  7. vLLM 部署 → OpenAI 兼容 API,一行切换

从训练到部署,32G 消费级显卡全部搞定。关键在于每一步的显存优化策略。

学AI大模型的正确顺序,千万不要搞错了

🤔2026年AI风口已来!各行各业的AI渗透肉眼可见,超多公司要么转型做AI相关产品,要么高薪挖AI技术人才,机遇直接摆在眼前!

有往AI方向发展,或者本身有后端编程基础的朋友,直接冲AI大模型应用开发转岗超合适!

就算暂时不打算转岗,了解大模型、RAG、Prompt、Agent这些热门概念,能上手做简单项目,也绝对是求职加分王🔋

在这里插入图片描述

📝给大家整理了超全最新的AI大模型应用开发学习清单和资料,手把手帮你快速入门!👇👇

学习路线:

✅大模型基础认知—大模型核心原理、发展历程、主流模型(GPT、文心一言等)特点解析
✅核心技术模块—RAG检索增强生成、Prompt工程实战、Agent智能体开发逻辑
✅开发基础能力—Python进阶、API接口调用、大模型开发框架(LangChain等)实操
✅应用场景开发—智能问答系统、企业知识库、AIGC内容生成工具、行业定制化大模型应用
✅项目落地流程—需求拆解、技术选型、模型调优、测试上线、运维迭代
✅面试求职冲刺—岗位JD解析、简历AI项目包装、高频面试题汇总、模拟面经

以上6大模块,看似清晰好上手,实则每个部分都有扎实的核心内容需要吃透!

我把大模型的学习全流程已经整理📚好了!抓住AI时代风口,轻松解锁职业新可能,希望大家都能把握机遇,实现薪资/职业跃迁~

这份完整版的大模型 AI 学习资料已经上传CSDN,朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费

在这里插入图片描述

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐