Veedubin/AttackLM

# AttackLM [](LICENSE) [](ATTRIBUTION.md) [](requirements.txt) [](https://huggingface.co/unsloth/Qwen2.5-Coder-3B-Instruct-bnb-4bit) ## 这是什么？ AttackLM 是一个完整的 pipeline，用于对小型语言模型进行 fine-tuning，使其 成为胜任的红队 / AI 安全助手。训练数据基于 **MITRE ATT&CK** 技术，源自公开授权的开源项目 （Atomic Red Team、MITRE Caldera、Metasploit、Sigma、Infection Monkey、RTA， 以及用于覆盖 AI 安全的 prompt injection 和越狱语料库）。 该 pipeline 摄取 10 个 MITRE 战术桶和 6 个专门桶 （orchestrator 路由、AI 模型攻击、安全工具），并生成一个 QLoRA LoRA adapter，你可以直接将其插在 `Qwen2.5-Coder-3B-Instruct` 之上。 它的独特之处在于： - **数据 pipeline 中不使用 LLM。** 每个训练对都是确定性地 从上游源中提取的 —— 没有幻觉内容，没有 API 成本。 - **内置坐标下降 HPO。** 扫描 `lora_r` (8→512) 和 `lora_dropout` (0→0.5)，并在最终训练前选出最佳结果。 - **兼容 16GB 到 128GB VRAM。** `--max-length 2048` 的 3B QLoRA 可以 装入 4080 SUPER。开启 packing 后，128GB 显卡可运行 70B+ 模型。 ## 数据来源归属 **所有训练数据均为公开授权开源项目的转换成果。** 我们不声明对任何技术、命令、模块 或规则拥有作者身份 —— 原作者保留这些权利。每个上游 repo、其 license 及其对 AttackLM 训练组合的贡献均记录在 [**`/ATTRIBUTION.md`**](ATTRIBUTION.md) 中，并在 [**`/NOTICE`**](NOTICE) 中进行了汇总。 完整的逐源映射： | Source | Pairs | License | Repository | |---|---:|---|---| | Atomic Red Team | 2,506 | MIT | [redcanaryco/atomic-red-team](https://github.com/redcanaryco/atomic-red-team) | | MITRE Caldera / Stockpile | 608 | Apache-2.0 | [mitre/stockpile](https://github.com/mitre/stockpile) | | Caldera plugins (arsenal/manx/access) | 56 | Apache-2.0 | [mitre/caldera](https://github.com/mitre/caldera) | | Metasploit Framework | 8,349 | BSD-3-Clause | [rapid7/metasploit-framework](https://github.com/rapid7/metasploit-framework) | | Infection Monkey | 36 | GPL-3.0 | [guardicore/monkey](https://github.com/guardicore/monkey) | | RTA — Red Team Automation | 76 | **AGPL-3.0** ⚠️ | [endgameinc/RTA](https://github.com/endgameinc/RTA) | | Sigma rules | (labels) | DRL-1.1 | [SigmaHQ/sigma](https://github.com/SigmaHQ/sigma) | | AI-security tools (promptfoo, garak, promptmap, PyRIT, FuzzyAI, TheBigPromptLibrary) | 743+ | mixed MIT/Apache-2.0 | various (see [ATTRIBUTION.md](ATTRIBUTION.md)) | | Synthetic orchestrator / prompt-injection | 1,067 | MIT | this repo | | **Total** | **16,982** | | | ⚠️ **AGPLv3 注意：** RTA 是唯一基于 AGPL 许可的源。AGPL 对 衍生作品具有网络分发影响。公开的 代码库满足了源代码可用性要求。如果你需要 纯净的 AGPL 部署，请在移除 `tools/rta` 桶后重新训练。 完整分析请参见 [ATTRIBUTION.md §8](ATTRIBUTION.md)。 ## 快速开始（5 分钟） ``` # 1. 安装 uv（Python 包管理器，约 10MB） curl -LsSf https://astral.sh/uv/install.sh | sh # 2. 克隆此 repo git clone https://github.com/Veedubin/AttackLM.git cd AttackLM # 3a. 作为 Python 包安装（为您提供 11 个 `attacklm-*` 命令） # — 使用 `[all]` 以获取所有可选依赖 uv pip install -e ".[all]" # 或者，如果您只想要基本的 CLI 调度器（无 ML stack）： # uv pip install -e . # 3b. 替代方案：使用经典的 uv 管理 venv 并包含 pyproject.toml 中的所有依赖 # uv sync # 4. 初始化数据集（优先探测本地 `data/`；若不存在则回退至 git clone） attacklm-init --yes # 上面的单条命令取代了下方的步骤 4–7。如果您更愿意 # 逐步单独运行，这四个命令依然可用： # # 4. 克隆上游数据源（总计约 1.5GB，可选 — 数据已包含在 repo 中） # attacklm-clone # # 5. 从每个来源中提取训练数据 # attacklm-extract # # 6. 为每个 JSONL 增加每对数据的来源/许可证归属信息 # attacklm-attribute # # 7. 组织分类到 16 个 MITRE/AI/tools 桶中 # attacklm-buckets # 8. 选择一个 base model — 使用 uncensored/abliterated 的模型（参见下方的“选择 base model”） # 示例：移除了 refusal direction 的 Qwen2.5-Coder-3B-Instruct # v0.2.0+ 版本使用 --dataset（多位置参数）替代了 --include-tools 等。 attacklm-train-all --single-model \ --dataset base/ \ --base-model huihui-ai/Qwen2.5-Coder-3B-Instruct-abliterated \ --epochs 5 --max-length 2048 # 可选：添加 --hho 以实现自动的 lora_r / lora_dropout sweep ``` 训练好的 LoRA adapter 会存放在 `models/attacklm-single_/` 中 （v0.2.0+ 使用带时间戳的目录，以便多次运行并存以供回滚）。 合并后的模型会输出到 `models/merged/attacklm-single/`。关于如何使用，请参阅 下方的 [**推理**](#inference)。 ## 安装 该项目以**标准的 Python 包**形式发布（`pyproject.toml`、 `src/attacklm/` 布局，hatchling 构建后端），因此用户无需 手动构建任何内容。 共有**两种 GPU 技术栈** —— 请根据你的硬件进行选择。 ### CUDA stack (NVIDIA) — 主要 ``` git clone https://github.com/Veedubin/AttackLM.git cd AttackLM uv pip install -e ".[all]" ``` 这将安装所有内容：`torch`（来自 PyPI 的 CUDA wheel）、`bitsandbytes`、 `transformers`、`peft`、`trl`，以及 C++ 扩展 `flash-attn`、 `causal-conv1d` 和 `flash-linear-attention`（用于 Qwen3-Next 及类似的 混合线性注意力模型）。 | Component | Where it comes from | |---|---| | `torch`, `torchvision` | PyPI (CUDA build, auto-selected) | | `bitsandbytes` | PyPI (CUDA wheels) | | `flash-attn` | Built from source via pip (~5 min) | | `causal-conv1d` | Pre-built wheel from PyPI | | `flash-linear-attention` | Pre-built wheel from PyPI | ### ROCm stack (AMD) — 例如 MI300X, RX 7900 XTX, Strix Halo ROCm PyTorch wheels **不在 PyPI 上** —— 你必须添加 PyTorch 的索引 URL。`bitsandbytes` 0.49+ 的 wheel **仅提供 CUDA .so 文件** (cuda118/120/121/122/126) —— 在 ROCm 上，请使用 `--no-deps` 安装 bitsandbytes 并进行验证，或者直接跳过它（FP8 路径不需要它）。C++ 扩展（`flash-attn`, `causal-conv1d`, `flash-linear-attention`）**没有 ROCm 支持** —— 但建模代码拥有纯 PyTorch 的后备方案（较慢但可用）。 **重要：该使用哪个 ROCm 版本？** PyTorch ROCm 索引为每个 频道发布不同的 `torch` 版本。本 repo 中的版本固定 （`torch==2.12.0`, `torchvision==0.27.0`）仅在 **rocm7.1 / rocm7.2** 频道上提供。旧频道（rocm6.x）最高仅支持 torch 2.5-2.9，将无法解析该固定版本。 ``` # 1. 从 rocm7.2 channel 安装 ROCm PyTorch（包含 torch 2.12.0） uv pip install --index-url https://download.pytorch.org/whl/rocm7.2 \ torch==2.12.0 torchvision==0.27.0 # 2. 安装带有 ROCm meta-group 的 AttackLM git clone https://github.com/Veedubin/AttackLM.git cd AttackLM uv pip install -e ".[all-rocm]" ``` 安装后，请验证： ``` python -c "import torch; print('torch:', torch.__version__, '— hip:', torch.version.hip)" # 应该会输出类似以下内容：torch: 2.12.0+rocm7.2 — hip: 7.2.XXXXX ``` `[all-rocm]` 即 `attacklm[train-rocm,extract,convert]` —— 它会拉取 `peft`、`trl`、`accelerate`、`bitsandbytes`，并且**不包含**仅限 CUDA 的 C++ 扩展。 | Component | Where it comes from | |---|---| | `torch`, `torchvision` | PyTorch ROCm index (`+rocm7.2` build) | | `bitsandbytes` | PyPI (CUDA-only — works for non-FP8 paths; uninstall for FP8) | | `flash-attn` | **Not installed** — sdpa fallback in QLoRA | | `causal-conv1d` | **Not installed** — pure-PyTorch fallback in Qwen3-Next modeling | | `flash-linear-attention` | **Not installed** — pure-PyTorch fallback | ### CPU / Apple Silicon（仅限推理） ``` git clone https://github.com/Veedubin/AttackLM.git cd AttackLM uv pip install -e ".[infer]" ``` 在 CPU/MPS 上进行训练在技术上是可行的，但会**极其缓慢**。 仅用于试运行或针对 prompts 运行预训练的 adapter。 请选择 `[all-cuda]` 或 `[all-rocm]` 进行实际训练。 ### 11 个控制台脚本入口点 所有安装路径都会为你提供以下命令： | Command | Dispatches to | What it does | |--------------------------|----------------------------------------|----------------------------------------| | `attacklm-train` | `scripts/train_template.py` | 训练单个 QLoRA adapter | | `attacklm-train-all` | `scripts/train_all.py` | 训练所有 buckets / HPO | | `attacklm-hpo` | `scripts/hpo_runner.py` | 坐标下降 HPO 扫描 | | `attacklm-infer` | `scripts/infer.py` | 冒烟测试推理 | | `attacklm-merge` | `scripts/merge_adapter.py` | 合并 LoRA → 基础模型 | | `attacklm-gguf` | `scripts/convert_to_gguf.py` | 转换为 GGUF (llama.cpp) | | `attacklm-build` | `scripts/build.py` | 合并 → GGUF → 安装（一步到位） | | `attacklm-demo` | `scripts/demo.py` | 多 agent orchestrator 演示 | | `attacklm-extract` | 所有 6 个提取器 | 从克隆的 repo 中提取数据 | | `attacklm-buckets` | `setup_buckets.py` + `reorganize_buckets.py` | 将数据整理为 16 个 buckets | | `attacklm-attribute` | `scripts/augment_attribution.py` | 为每个 JSONL 行添加来源/ license | | `attacklm-clone` | `scripts/clone_repos.sh` | 克隆上游数据 repo | | `attacklm-init` | `scripts/init_pipeline.py` | **一次性初始化：克隆→提取→归属→分桶** (优先探测本地) | | `attacklm-balance` | `scripts/balance_buckets.py` | 构建 buckets 的平衡子集 | CLI 分发器是轻量级的包装器 —— 它们使用 `runpy.run_path()` 来调用 `scripts/` 中的标准脚本。因此 `scripts/` 仍然是 事实来源，如果你愿意，你仍然可以直接运行 `uv run python scripts/foo.py`。 ### 可选依赖组（高级） ``` # 细粒度控制 uv pip install -e ".[train-cuda]" # CUDA training stack uv pip install -e ".[train-rocm]" # ROCm training stack uv pip install -e ".[infer-cuda]" # CUDA inference uv pip install -e ".[infer-rocm]" # ROCm inference uv pip install -e ".[extract]" # data extractors uv pip install -e ".[convert]" # GGUF conversion uv pip install -e ".[dev]" # pytest, ruff, mypy ``` ### 免安装选项（仅脚本） 如果你不想将其安装到你的环境中： ``` git clone https://github.com/Veedubin/AttackLM.git cd AttackLM uv sync # creates .venv with all deps uv run python scripts/train_all.py --single-model --epochs 5 ``` `uv sync` 会读取 `pyproject.toml` 并创建一个带有 `[all]` extras 的虚拟环境。`scripts/` 中的脚本是事实来源 —— CLI 只是一个 轻量级的分发层。 ## 架构 训练数据被组织为 **16 个 buckets**： - **10 个 MITRE 战术 buckets** —— 位于 `base/` 下：`base/collection`、 `base/command_and_control`、`base/credential_access`、`base/defense_evasion`、 `base/discovery`、`base/execution`、`base/exfiltration`、 `base/lateral_movement`、`base/persistence`、`base/privilege_escalation` （分别对应 TA0009、TA0011、TA0006、TA0005、TA0007、TA0002、TA0010、TA0008、 TA0003、TA0004） - **1 个 orchestrator bucket** —— 跨 6 个子 agent 的路由决策 - **2 个 AI 模型攻击 buckets** —— 位于 `ai/` 下：`ai/prompt-injection` 和 `ai/jailbreaking`（TA0040 — 对抗性 ML） - **3 个安全工具 buckets** —— 位于 `tools/` 下：`tools/infection_monkey`、 `tools/metasploit`、`tools/rta`（汇总了特定工具的数据，在适用的情况下重新路由 到 MITRE 战术） bucket 的布局让你可以训练： - **在一个包含所有数据的模型上**（默认 —— 单个 MoE 风格的助手） - **每种战术一个模型**（多模型模式） - **在子集上训练一个模型**（例如，`--include-tools --include-orchestrator` 以跳过 AI/ML 攻击数据） 请查看 `data/datasets/buckets/manifest.json` 以获取每个 bucket 的完整清单， 包括对数计数和 MITRE 战术 ID。 ## 选择基础模型 ### 推荐基础模型（任选其一） | Model | Size | VRAM needed | Notes | | ------------------------------------------------------------ | ----- | ------------ | -------------------------------------------------------------------- | | `huihui-ai/Qwen2.5-Coder-3B-Instruct-abliterated | 3B | 16 GB | 最适合 RTX 4080 SUPER / 4070 Ti。与原版默认配置采用相同的 Qwen2.5-Coder 架构。Apache-2.0。 | | `huihui-ai/Qwen2.5-Coder-1.5B-Instruct-abliterated` | 1.5B | 8 GB | 硬件吃紧，迭代快。Apache-2.0。 | | `huihui-ai/Qwen2.5-Coder-7B-Instruct-abliterated` | 7B | 24 GB | 质量更好，长回复更连贯。Apache-2.0。 | | `BlossomsAI/Qwen2.5-Coder-32B-Instruct-Uncensored` | 32B | 64+ GB | 顶级质量，需要 64+ GB VRAM。Apache-2.0。 | | `failspy/Meta-Llama-3-8B-Instruct-abliterated-v3` | 8B | 24 GB | 如果你更想在 Llama-3 上进行训练。Apache 风格 license。 | | `failspy/Meta-Llama-3-70B-Instruct-abliterated-v3.5` | 70B | 128+ GB | 前沿质量。还提供量化后的 GGUF 版本。 | 浏览完整的 [failspy/abliterated-v3 collection](https://huggingface.co/collections/failspy/abliterated-v3-664a8ad0db255eefa7d0012b) 和 [3000+ Heretic models](https://huggingface.co/models?other=heretic) 了解更多。 ### 使用 Heretic 制作你自己的（如果你首选的基础模型未经过 abliteration 处理） [p-e-w/heretic](https://github.com/p-e-w/heretic) 是一个全自动的 abliteration 工具。在 16 GB 显卡上为 3B 模型处理只需 30 分钟。 ``` pip install heretic-llm heretic Qwen/Qwen2.5-Coder-3B-Instruct --n-trials 100 # 交互式菜单：选择“Save the model to a local folder” ``` 然后将 `--base-model` 指向保存的文件夹。交互式菜单 **需要一个真实的 TTY**（gnome-terminal、konsole、xterm 等）—— 通过 `printf "n\n"` 管道传输 stdin 仅能处理第一个提示。 该技术其余的 30 行代码记录在以下文档中： - [mlabonne/abliteration（2024 年原始方案）](https://huggingface.co/blog/mlabonne/abliteration) - [grimjim/projected-abliteration（2025 年 10 月 —— 投影细化）](https://huggingface.co/blog/grimjim/projected-abliteration) - [p-e-w/heretic（统一工具，现代最先进方案）](https://github.com/p-e-w/heretic) ## 训练 `scripts/train_all.py` 是 orchestrator。关键 flags： | Flag | Default | Notes | |---|---|---| | `--single-model` | (off) | 在所有合并后的 buckets 上训练一个模型 | | `--base-model` | (auto) | v0.2.0+：默认为第二轮 SFT（该 agent 最近完成的一次运行），然后是 abliterated 的 Qwen 3B。传递此参数以覆盖。 | | `--dataset` (multi) | none | v0.2.0+：bucket 规格的位置参数列表。`base/`, `tools/`, `ai/`, `orchestrator`, 子路径 (`tools/metasploit/`), 别名 (`all`, `tactics`, `tools-all`). | | `--backup` | (on) | 在训练开始前，将上一轮的 round-2 SFT 运行 Tar.gz 打包到 `models/.backups/`。使用 `--no-backup` 跳过。 | | `--epochs` | 10 | 合并数据集上的总 epoch 数 | | `--max-length` | 1024 | 2048 用于更丰富的上下文；1024 用于在 16GB 上的 7B 模型 | | `--lora-r` | 16 | LoRA 秩；8 / 16 / 32 是不错的起始点 | | `--lora-alpha` | 32 | 通常为 `2 × lora_r` | | `--lora-dropout` | 0.05 | 尝试设为 0.0 以减少正则化 | | `--no-packing` | (packing off) | 默认为 OFF，因为 flash-attn 很难安装 | | `--packing` | (off) | 启用可提升约 30% 速度；需要 `flash_attn` | | `--include-tools` | (off) | **在 v0.2.0 中已弃用**：改用 `--dataset tools/` | | `--include-orchestrator` | (off) | **在 v0.2.0 中已弃用**：改用 `--dataset orchestrator` | | `--model-attacks` | (off) | **在 v0.2.0 中已弃用**：改用 `--dataset ai/` | | `--curriculum` | (off) | 2 阶段：先训练战术数据，然后微调 orchestrator | | `--hpo` | (off) | 在最终训练前运行坐标下降 HPO | 训练脚本内置了 13 项 OOM 安全修复（expandable_segments、 per_device_eval_batch_size=1、chunked_nll loss、post-eval cache clear、 paged_adamw_8bit 等）—— 请查看 `train_template.py` 中的 `# OOM fix #N:` 注释以获取完整列表。 ### 运行目录命名（v0.2.2+） `attacklm-train` 和 `attacklm-train-all` 默认都会将 adapter 写入一个**带时间戳的**子目录中，以便保留重复运行的结果： ``` # 默认行为 — 在您的 --output 后附加时间戳 attacklm-train --dataset data/foo.jsonl --output models/agent-3b # → models/agent-3b_2026-06-10_15-15/ （在重新运行时保留） # 选择停用时间戳（将拒绝覆盖已完成的运行） attacklm-train --dataset data/foo.jsonl --output models/agent-3b --no-timestamp # ERROR: Refusing to clobber completed run at models/agent-3b. # 强制覆盖拒绝限制 attacklm-train --dataset data/foo.jsonl --output models/agent-3b \ --no-timestamp --force ``` 如果 `--output` 已经以 `_YYYY-MM-DD_HH-MM` 结尾（即它是 由之前的运行或由 `attacklm-train-all` 生成的），则该后缀 将保持不变 —— 仅当存在完全 相同的名称时，重复运行才会获得新的后缀（`_2`, `_3`, …）。 ### 多轮 SFT（v0.2.0+） 每次训练运行都会在 `models/{agent}_{TIMESTAMP}/state.json` 写入一个 `state.json` 附带文件。 它会记录基础模型、hparams、数据集、进度以及一个 `completed` 标志。 **第二轮 SFT** 会在之前完成的运行之上训练一个新的 LoRA： ``` # 第 1 轮：在 tactics 上训练（10 个桶，7,398 对） attacklm-train-all --single-model --dataset base/ --epochs 5 # 第 2 轮：在第 1 轮合并后的权重之上训练 tools # （从 state.json 自动检测；首先会进行第 1 轮的备份 tar） attacklm-train-all --single-model --dataset tools/ --epochs 3 # 第 3 轮：在所有内容上训练 attacklm-train-all --single-model --dataset all --epochs 2 ``` 每一轮： 1. 检测该 agent 名称最近完成的运行 2. 将其备份到 `models/.backups/{name}_{timestamp}.tar.gz`（5 GB，约 30 秒） 3. 加载合并后的权重作为新的基础模型 4. 在此之上训练一个新的 LoRA 5. 写入一个新的带有时间戳的运行目录以及更新后的 `state.json` 针对崩溃/被终止运行的**自动恢复**： ``` # 如果运行在训练过程中中断，只需使用相同的命令重新运行即可。 # state.json (completed=false) + checkpoint-N/ 存在 → 自动恢复。 attacklm-train-all --single-model --dataset base/ --epochs 5 ``` ### `--dataset` DSL 新的数据集规格采用目录形状的分层结构： | Spec | Resolves to | Pair count | |-------------------------------|------------------------------------------------------|-----------:| | `base/` | All 10 MITRE tactic buckets | 7,398 | | `tools/` | All 3 tool buckets (metasploit, infection_monkey, rta) | 8,461 | | `tools/metasploit/` | Just metasploit | 8,349 | | `tools/infection_monkey/` | Just infection_monkey | 36 | | `tools/rta/` | Just RTA | 76 | | `ai/` | Both AI buckets (jailbreaking, prompt-injection) | 743 | | `orchestrator` | The orchestrator bucket | 380 | | `all` | Everything (alias for `base + tools + ai + orchestrator`) | 16,982 | | `tactics` | Alias for `base/` | 7,398 | 多个规格可组合使用：`--dataset base/ tools/metasploit/` = 10 个战术 + 仅 metasploit = 15,747 对。 原有的 `--include-tools` / `--model-attacks` / `--include-orchestrator` 仍然有效， 并会在内部转换为 `--dataset` 规格。如果同时传递，新 flag 优先。 ### 平衡采样（`attacklm-balance`） 这 16 个 buckets 的数据量严重倾斜：仅 `tools/metasploit` 就有 8,349 对（占 16,982 总数的 49%）。直接在原始的 `--dataset all` 上训练 会使模型每看到 1 个非 Metasploit 示例，就会看到约 2 个 Metasploit 示例， 这会导致其对 msfconsole 语法过拟合，从而牺牲了 更广泛的战术覆盖。 `attacklm-balance` 构建了一个平衡的 buckets 子集。它会应用 一个每个 bucket 的上限（一个统一应用于所有 buckets 的上限），并 使用选定的策略从每个 bucket 中选择示例： ``` # 试运行：查看每个桶的上限 + 总数，而不进行写入 attacklm-balance --profile 7b-128gb --dry-run # 将平衡后的 dataset 写入 data/datasets/balanced/ attacklm-balance --profile 7b-128gb \ --output data/datasets/balanced/balanced_7b-128gb.jsonl # 然后基于此进行训练 attacklm-train --dataset data/datasets/balanced/balanced_7b-128gb.jsonl \ --output models/attacklm-7b-128gb \ --base-model huihui-ai/Qwen2.5-Coder-7B-Instruct-abliterated ``` **配置文件**（每个 bucket 的命名上限值，针对常见硬件组合进行了调整）： | Profile | Per-bucket cap | Total pairs | Notes | |------------|---------------:|------------:|--------------------------------------------| | `3b-16gb` | 800 | ~7,500 | 3B QLoRA on 16 GB card | | `7b-16gb` | 800 | ~7,500 | 7B QLoRA on 16 GB card | | `7b-128gb` | 1,500 | ~9,800 | 7B QLoRA on 128 GB rig | | `14b-128gb`| 1,500 | ~9,800 | 14B QLoRA on 128 GB rig | | `31b-128gb`| 2,000 | ~10,600 | 31B QLoRA on 128 GB rig | | `full` | unlimited | 16,982 | All data, no cap | | `custom` | (you set)| (you set)| `--per-bucket-cap` or `--target-total` | **策略**（在 bucket 内，应用上限后）： - `stratified` (default) —— 根据示例的首个 MITRE 技术 ID、来源或助手内容的第一行进行分组，然后 分配**每组至少 1 个**示例，以确保每种技术 / 模块都有 代表。如果 bucket 中少于 3 个组，则回退为均匀随机。 - `random` —— N 的均匀随机样本（由 `--seed` 设置种子）。 - `head` —— 文件中的前 N 个示例（可复现，但偏向于 数据所处的任何顺序）。 **自定义分配** —— `custom` 配置文件接受显式的 `--per-bucket-cap` JSON 或带有 `--category-shares` 的 `--target-total`： ``` # 总计 12K 对数据，按 30% tactics / 40% tools / 20% ai / 10% orchestrator 加权 attacklm-balance --profile custom --target-total 12000 \ --category-shares '{"tactic": 0.3, "tools": 0.4, "ai_redteam": 0.2, "meta": 0.1}' # 仅设置 metasploit 为 1500 + discovery 为 800，其余不限 attacklm-balance --profile custom \ --per-bucket-cap '{"tools/metasploit": 1500, "base/discovery": 800}' ``` 输出的 JSONL 会写入到 `data/datasets/balanced/`，被排除在 git 之外， 并在每个示例中包含一个 `_source_bucket` 字段以便于追踪。 请查看 `scripts/balance_buckets.py --help` 获取完整的 flag 列表，并查看 `CHANGELOG.md` 了解设计原理。 ### HPO 将 `--hpo` 添加到训练命令中。该扫描会探索 `lora_r` (8→512) 和 `lora_dropout` (0→0.5)，并使用获胜者进行最终训练。 结果将存放在 `hpo_runs/hpo_state.json` 中；以后可以通过 `attacklm-hpo --analyze-only` 重新分析。 ## 推理 训练后，你会在 `models/attacklm-single_*/`（带时间戳）中得到一个或多个 LoRA adapter。 挑选最新的一个（最近日期）并将其合并。有三种使用方式： ### 选项 A：使用 `infer.py` 进行快速冒烟测试 ``` # v0.2.0+：列出可用的 run dirs 并选择最新的一个 ls -d models/attacklm-single_*/ | tail -1 # 然后针对其进行推理 attacklm-infer --adapter models/attacklm-single_2026-06-10_01-12 ``` 这会运行 4 个示例 prompt（MITRE 战术、orchestrator 路由、 prompt injection）并打印模型的回复。除了 `uv sync` 之外无需任何设置。请查看 `scripts/infer.py --help` 以使用自定义 prompt 和生成参数。 ### 选项 B：合并到基础模型中（最简单） ``` # v0.2.0+：--adapter 直接接受一个带时间戳的 run dir。 # 如果您省略 --adapter，merge_all 会自动为 agent 选择最新的运行。 attacklm-merge \ --base-model huihui-ai/Qwen2.5-Coder-3B-Instruct-abliterated \ --adapter models/attacklm-single_2026-06-10_01-12 \ --output models/merged/attacklm-single ``` 然后使用 `transformers.AutoModelForCausalLM.from_pretrained("models/merged/attacklm-single")` 加载。 ### 选项 C：转换为 GGUF 以供 Ollama / LM Studio / llama.cpp 使用 ``` # v0.2.0+：--input 是合并后的 model dir（不是 adapter） attacklm-gguf \ --input models/merged/attacklm-single \ --install-lmstudio # 注册到 Ollama uv run python scripts/register_ollama.py models/gguf/attacklm-single.Q4_K_M.gguf ``` ### 选项 E：一键合并 + GGUF + 安装（`attacklm-build`） v0.2.2+：3 条命令的 shell pipeline 变成了单个命令。该 构建命令还会在 `models/built/{name}_{timestamp}/` 生成一个清单 以供日后检索： ``` # 合并 + GGUF + 安装到 LM Studio，一步到位 attacklm-build \ --adapter models/attacklm-3b_16g_2026-06-10_15-15 \ --base ./uncensored/ \ --name attacklm-3b-16g # 跳过合并步骤（使用已合并的 model） attacklm-build \ --merged models/merged/attacklm-3b-16g \ --name attacklm-3b-16g # 同时注册到 Ollama attacklm-build \ --adapter models/attacklm-3b_16g_2026-06-10_15-15 \ --base ./uncensored/ \ --name attacklm-3b-16g \ --register-ollama ``` `install-lmstudio` 默认开启。使用 `--no-install-lmstudio` 以 仅生成 GGUF。构建清单记录了 GGUF 路径、 mtime、基础模型以及执行了哪些安装步骤。 ### 选项 D：直接加载 adapter（占用磁盘空间最小） ``` from peft import PeftModel from transformers import AutoModelForCausalLM, AutoTokenizer base = AutoModelForCausalLM.from_pretrained( "huihui-ai/Qwen2.5-Coder-3B-Instruct-abliterated", device_map="auto", ) model = PeftModel.from_pretrained(base, "models/attacklm-single") tokenizer = AutoTokenizer.from_pretrained("models/attacklm-single") # 与模型进行聊天 messages = [ {"role": "system", "content": "You are an authorized Red Team specialist..."}, {"role": "user", "content": "Show the System Services: Service Execution technique (T1569.002)"}, ] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(text, return_tensors="pt").to(model.device) outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7) print(tokenizer.decode(outputs[0], skip_special_tokens=True)) ``` ## 数据来源（上游） | Project | License | Use | |---|---|---| | [redcanaryco/atomic-red-team](https://github.com/redcanaryco/atomic-red-team) | MIT | 2,506 atomic test triples | | [mitre/stockpile](https://github.com/mitre/stockpile) | Apache-2.0 | 608 adversary-emulation abilities | | [mitre/caldera](https://github.com/mitre/caldera) | Apache-2.0 | 56 plugin descriptors | | [rapid7/metasploit-framework](https://github.com/rapid7/metasploit-framework) | BSD-3-Clause | 8,349 module description triples | | [guardicore/monkey](https://github.com/guardicore/monkey) | GPL-3.0 | 36 plugin manifest triples | | [endgameinc/RTA](https://github.com/endgameinc/RTA) | **AGPL-3.0** ⚠️ | 76 Python TTP triples | | [SigmaHQ/sigma](https://github.com/SigmaHQ/sigma) | DRL-1.1 | Auxiliary context for triple structure | | [promptfoo/promptfoo](https://github.com/promptfoo/promptfoo) | MIT | Prompt injection probes | | [NVIDIA/garak](https://github.com/NVIDIA/garak) | Apache-2.0 | DAN/probe resources | | [utkusen/promptmap](https://github.com/utkusen/promptmap) | MIT | Prompt injection rules | | [Azure/PyRIT](https://github.com/Azure/PyRIT) | MIT | Jailbreak templates | | [cyberark/FuzzyAI](https://github.com/cyberark/FuzzyAI) | Apache-2.0 | Adversarial prompt resources | | [Resident-Falker/TheBigPromptLibrary](https://github.com/Resident-Falker/TheBigPromptLibrary) | mixed MIT/MPL | Jailbreak + system prompt library | 完整的归属说明、每对数据的来源映射以及再分发指南请参见 [**`/ATTRIBUTION.md`**](ATTRIBUTION.md)。 ## 许可证 - **本 repo 中的代码** —— [MIT License](LICENSE) - **训练数据** —— 继承其组件中最严格的 license （目前为来自 RTA 的 AGPL-3.0 —— 参见 [ATTRIBUTION.md §8](ATTRIBUTION.md)） - **训练后的模型权重** —— 作为从公开授权材料中学习到的 全新统计产物，采用 MIT License。模型权重在 版权意义上是否属于“衍生作品”是一个悬而未决的问题；我们 对此不作任何声明。如果你需要确定性，请就你的 具体部署场景咨询法律顾问。 上游 MITRE、NVIDIA 和 CyberArk 组件所要求的 Apache-2.0 归属说明保留在 [**`/NOTICE`**](NOTICE) 中。 ## 贡献 有关提交 issue、提交 PR 以及扩展 bucket/extractor 系统的指南，请参见 [CONTRIBUTING.md](CONTRIBUTING.md)。 ## 更新日志 完整的版本历史请参见 [CHANGELOG.md](CHANGELOG.md)。近期值得注意的 发布版本： - **v0.2.2** (2026-06-10) —— `attacklm-balance`（平衡 bucket 采样器）， `attacklm-build`（一键 merge+GGUF+install），`attacklm-train` 中 自动添加时间戳的运行目录，准确的 epoch 计数器，基于 GGUF mtime 的陈旧检查，`attacklm-gguf --name` / `--register-ollama` / `--quant` / `--build` / `--force`。 - **v0.2.1** (2026-06-10) —— Bucket 布局标准化为 4 个父级 （`base/`, `tools/`, `ai/`, `orchestrator/`）。 - **v0.2.0** (2026-06-10) —— 多轮 SFT，`state.json` 溯源， `--dataset` DSL，`--backup`/`--no-backup`，GGUF 转换中的 LoRA adapter 检测。**主版本号提升。** - v0.1.5 (2026-06-10) —— LM Studio 路径修复，内核 pin，路径解析器 - v0.1.4 (2026-06-10) —— 合并 + GGUF pipeline - v0.1.0 (2026-06-10) —— 首次公开发布 ## 致谢 感谢开源安全社区 —— Red Canary、MITRE、Rapid7、 Guardicore、Endgame/Elastic、SigmaHQ 维护者、promptfoo、 garak、PyRIT 和 FuzzyAI 团队，以及所有为我们所依赖的项目 做出贡献的人。AttackLM 站在他们的肩膀上。

标签：DLL 劫持, QLoRA, 人工智能, 凭据扫描, 大语言模型, 安全攻防, 数据展示, 模型微调, 用户模式Hook绕过, 知识库安全, 系统调用监控, 红队, 训练数据集, 逆向工具