BorisTerminator/llm-guard-lora
GitHub: BorisTerminator/llm-guard-lora
一个基于 QLoRA 微调的轻量级 prompt-injection guard 模型,通过独立的分类器对用户请求进行三级安全判定,在不显著增加延迟的前提下隔离越狱攻击面。
Stars: 0 | Forks: 0
# llm-guard-lora
一个 QLoRA 微调模型,能将 3B 的 instruct 模型转化为 guard 分类器:它会读取传入的用户请求,并决定其背后的 assistant 是可以回答、必须在限制条件下回答,还是必须拒绝。
它的体积足够小巧,可以部署在每个请求的前端,而不会增加明显的延迟或成本。
## 为什么需要微调而不是 prompt
作为主模型 system prompt 实现的 guard 会共享该模型的 context,而任何共享 context 的内容都是可以被辩驳的。将决策拆分到一个只输出判定结果的独立小模型中,可以消除这种攻击面:这里没有可以被劫持的对话,针对 assistant 的越狱攻击永远无法触达决定 assistant 是否应该回复的组件。
这种分离的代价是在路径中增加了一个额外的模型,这就是它必须足够小的原因。
## 分类体系
三种判定结果下共有十七个类别。中间层是此设计的核心所在——大多数真实请求既非完全安全,也非绝对禁止,而是可以在限制条件下回答的。
| 判定 | 代码 | 含义 |
|---|---|---|
| **GREEN** | `G1` | 正常回答 |
| **YELLOW** | `Y1`–`Y6` | 回答,但需在明确限制下 |
| **RED** | `B1`–`B10` | 拒绝 |
每个类别都带有一个机器可读的 `reason`,调用应用程序会将其注入到 assistant 自身的 prompt 中,因此 guard 不仅仅起到门控作用——它还提供指导。
黄色限制示例:
| 代码 | 情境 | 返回的限制 |
|---|---|---|
| `Y2` | 一般医疗问题 | 不提供剂量、诊断或治疗方案;建议咨询医生 |
| `Y3` | 供应商选择 | 不推荐特定产品或公司;转而协助定义选择标准 |
| `Y4` | 个人理财 | 不提供投资建议;建议咨询专家 |
| `Y6` | 大量数据导出 | 建议按筛选条件、时间段和结果数量进行缩小范围 |
红色类别涵盖了对非法或危险指令、医疗诊断、政治劝导及相关情况的需求。
## 训练设置
| | |
|---|---|
| 基础模型 | `Qwen/Qwen2.5-3B-Instruct`(也支持 1.5B 和 7B) |
| Quantisation | 4-bit NF4,double quantisation,bfloat16 计算 |
| LoRA rank / alpha | 16 / 32,dropout 0.05 |
| 目标模块 | `q_proj` `k_proj` `v_proj` `o_proj` `gate_proj` `up_proj` `down_proj` |
| Epochs | 3 |
| 有效批次大小 | 16(batch 2 × gradient accumulation 8) |
| 学习率 | 2e-4,余弦调度,5% warmup |
| 权重衰减 | 0.01 |
| 检查点选择 | 最佳 `eval_loss`,每轮 epoch 评估一次 |
在这里,调整所有投影层而不仅仅是 attention 层非常重要:该任务是以生成形式表达的分类,而 MLP 块承载了大部分类别区分功能。
## 评估
`train_lora.py --eval-only` 会从三个维度对已保存的 adapter 进行评分,因为它们失败的方式各不相同:
- **类别准确率** — 是否选出了确切的代码
- **判定准确率** — 是否选对了层级,即使其中的代码不对
- **违规准确率** — 是否正确决定拦截还是放行
将 `Y2` 预测为 `Y4` 是一个轻微错误;而将 `B3` 预测为 `G1` 则是整个系统的失败。将这三者区分开来可以清晰地展现这一点,而不是将其平均化掩盖掉。运行结果还会打印每个类别的准确率和判定混淆矩阵。
## 仓库结构
| 路径 | 用途 |
|---|---|
| `guard_prompt.txt` | 定义分类体系的 system prompt,包含 `{{user_input}}` 和 `{{language}}` 占位符 |
| `generate_dataset.py` | 生成合成的带标签 guard 示例 |
| `train_lora.py` | QLoRA 训练和评估 |
| `make_datasets.ipynb` | 数据集组装与拆分 |
| `help.ipynb` | 探索与检查 |
| `Datasets/` | 训练集和保留的测试集 |
| `references.md` | 调研过的 prompt 注入模型和数据集 |
## 运行方式
```
pip install transformers peft trl datasets bitsandbytes accelerate pandas openpyxl
python train_lora.py # train and evaluate
python train_lora.py --eval-only # evaluate a saved adapter
python train_lora.py --model Qwen/Qwen2.5-1.5B-Instruct # smaller base model
```
## 技术栈
PyTorch · Transformers · PEFT · TRL · bitsandbytes · Qwen2.5
标签:DLL 劫持, LoRA微调, 人工智能, 内容分类, 凭据扫描, 大语言模型, 提示注入防御, 模型网关, 源代码安全, 用户模式Hook绕过, 系统调用监控, 逆向工具