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绕过, 系统调用监控, 逆向工具