RII6/LLM-Safety-platform
GitHub: RII6/LLM-Safety-platform
通过分析 LLM 内部状态（logits 和残差流激活）而非仅依赖输出文本，为微调后的模型 checkpoint 提供部署前安全审计与分级判定。
Stars: 6 | Forks: 0
# 内部状态 LLM 安全扫描器
一个专为 **ML 工程师** 设计的微调 LLM checkpoint 部署审计工具，
适用于那些微调了模型并希望得到“是否可以安全发布？”这一答案的人。
与发送 prompt 并对**响应文本**进行分类的文本级红队工具不同，该扫描器读取模型的**内部状态**——输出 logits 和残差流激活（residual-stream activations）。一个模型可以打印出礼貌的拒绝，而其 logits 早已准备好服从，文本分析会漏掉这一点，但内部状态不会。
## 它报告的内容
对于每个 checkpoint，扫描器会输出一个包含诊断和建议的判定（verdict）：
| 判定 | 含义 |
| -------------------------- | ----------------------------------------------------------------- |
| `deploy` | 拒绝危害，在内部表示出危害 —— 健康 |
| `deploy_with_restrictions` | 部分服从 —— 仅在带有 guardrails 的情况下发布 |
| `do_not_deploy` | 在有害 prompt 上服从 |
| `inconclusive` | 行为探针未携带安全信号（例如 base model） |
有关判定背后的指标和设计决策，请参阅 [`safety_docs.md`](docs/safety_docs.md)。
## 模块
- **`safety_margin`** — 行为信号。Teacher-forces 规范化的拒绝/服从续写，并计算 `margin = logP(refuse) − logP(comply)`；`fail_rate` 是指在有害 prompt 中，服从成为更低成本路径的比例。这是 GCG 攻击目标（Zou 等人，2023），因此它是一个经过验证的越狱能力（jailbreakability）代理指标。
- **`refusal_direction`** — 内部状态信号。残差流中的均值差异“拒绝方向”（Arditi 等人，2024）；有害数据与良性数据分离程度的 AUROC 衡量了对危害的**感知**能力。
- **`verdict`** — 将行为 × 感知能力聚合为部署决策。
## 设置
需要 [uv](https://docs.astral.sh/uv/)。可在 **CUDA（NVIDIA，包括 Windows）、Apple Silicon (MPS) 或 CPU** 上运行 —— 设备是自动检测的（CUDA > MPS > CPU），因此同一个命令可以在团队的不同机器上运行。
```
uv sync
```
在带有 NVIDIA GPU 的 Windows 上，默认的 PyPI `torch` 仅支持 CPU。要使用 GPU，请安装与你的驱动程序匹配的 CUDA 版本，例如：
```
uv pip install torch --index-url https://download.pytorch.org/whl/cu124
```
## 构建语料库
语料库（`src/data/corpus/{harmful,benign}.jsonl`）已预构建。要从公开数据集（AdvBench 有害行为 + Alpaca 指令）重新生成或调整其大小：
```
uv run python scripts/build_corpus.py --n 500
```
这两个类别都是祈使句风格，这控制了 prompt 格式的混淆因素。
## 运行
```
uv run python scripts/main.py # full corpus (500/500), full run is slow on CPU
uv run python scripts/main.py --sample 30 # fast dev run, 30 prompts per class
uv run python scripts/main.py --device cpu # force a backend (default: auto-detect)
```
`scripts/main.py` 运行三个真实标准（ground-truth）的 checkpoint —— 一个对齐模型、一个消融（abliterated）模型（手术式移除了安全性），以及一个 base 模型（从未进行过安全微调）—— 并打印每个模块的摘要以及判定。一个正确的扫描器必须能够将它们区分开。
Windows 注意事项：NVIDIA GPU 会自动被识别为 CUDA；如果没有 GPU，它会回退到 CPU（正确，只是更慢 —— 使用 `--sample` 进行迭代）。
## Web 应用
一个单页 UI：粘贴一个 Hugging Face repo，即可获得判定以及每个指标的通俗语言解释。扫描是实时运行的，因此目标 VM 的模型大小有上限。
```
uv run uvicorn src.app.server:app --reload --port 8000 # http://localhost:8000
```
旋钮（环境变量）：`SCAN_SAMPLE`（每类 prompt 数，默认 25），`SCAN_MAX_PARAMS`（在下载前拒绝较大的模型，默认 400M），`SCAN_DTYPE`（默认 `bfloat16`），`SCAN_DEVICE`（默认 `cpu`）。一次仅运行一个扫描；报告会缓存到 `reports/`。
## Docker（推荐）
运行整个应用程序（前端、后端和 PostgreSQL）最简单的方法是使用 Docker Compose。它会构建前端、设置 Python 环境并为你运行数据库。
```
docker compose up -d
```
Web 应用程序将通过 `http://localhost` 访问。模型缓存在 Docker volume 中，因此它们不会在重启时重新下载。
## 部署（2 GB VM）
大小适合约 1.5 GB 的可用 RAM。torch (CPU) 约 0.5 GB + FastAPI 约 0.12 GB，为权重留下了其余空间，因此该应用程序限制了模型大小并序列化扫描。
```
# swap 作为安全网
sudo fallocate -l 2G /swapfile && sudo chmod 600 /swapfile
sudo mkswap /swapfile && sudo swapon /swapfile
# install (仅 CPU 的 torch 保持其精简)
uv venv && uv pip install torch --index-url https://download.pytorch.org/whl/cpu && uv sync
# 在 systemd 下运行 (MemoryMax 限制失控的 scan)
sudo cp deploy/capstone.service /etc/systemd/system/
sudo systemctl enable --now capstone
```
`deploy/capstone.service` 设置了 `MemoryMax=1500M`。在前面放置 nginx/caddy 以实现 TLS，或直接暴露 80 端口以实现简单的 MVP。
标签：AI红队, DLL 劫持, GraphQL安全矩阵, Vectored Exception Handling, 内部状态分析, 凭据扫描, 大语言模型, 安全扫描器, 模型安全评估, 测试用例, 版权保护, 逆向工具