142 lines
5.6 KiB
Markdown
142 lines
5.6 KiB
Markdown
# vllm-k100-patch-v2
|
||
|
||
## 背景
|
||
|
||
在海光K100-AI平台跑大批量模型 benchmark 时，部分模型因 tokenizer 配置问题或非标准 attention head_size 导致启动失败。本仓库通过在基础镜像上打 patch，修复已知的兼容性问题，无需逐个修改模型文件。
|
||
|
||
## 目录结构
|
||
|
||
```text
|
||
├── Dockerfile # 镜像构建文件
|
||
├── patch.py # transformers 库源码修复（构建时执行）
|
||
├── patch_triton.py # Triton backend head_size 验证绕过（构建时执行）
|
||
├── detect_tokenizer.py # 检测 tokenizer 类型（运行时调用）
|
||
├── fix_tokenizer.py # 修复 tokenizer 配置（运行时执行）
|
||
├── detect_head_size.py # 检测模型 head_size（运行时执行）
|
||
└── entrypoint.sh # 容器入口脚本，协调修复流程并启动 vLLM
|
||
```
|
||
|
||
## 镜像信息
|
||
|
||
| 项目 | 值 |
|
||
|------|----|
|
||
| 基础镜像 | `git.modelhub.org.cn:9443/enginex-hygon/vllm:0.9.2` |
|
||
| Patched 镜像（仅 tokenizer 修复）| `git.modelhub.org.cn:9443/enginex-hygon/vllm:0.9.2-patch-tokenizer` |
|
||
| Patched 镜像（tokenizer + head_size 修复）| `git.modelhub.org.cn:9443/enginex-hygon/vllm:0.9.2-patch-v2` |
|
||
|
||
## 已修复问题
|
||
|
||
### Patch 1：extra_special_tokens 类型兼容
|
||
|
||
**问题：** 部分模型的 `tokenizer_config.json` 中 `extra_special_tokens` 字段为 `list` 类型，而 transformers 库要求其为 `dict`，导致 vLLM 启动时报错崩溃。
|
||
|
||
**修复文件：** `transformers/tokenization_utils_base.py`
|
||
|
||
**修复方式：** 在调用 `.keys()` 之前检测类型，若为 `list` 则自动转换为 `dict`。
|
||
|
||
```python
|
||
if isinstance(special_tokens, list):
|
||
special_tokens = {t: t for t in special_tokens}
|
||
```
|
||
|
||
### Patch 2：Tokenizer 配置自动修复
|
||
|
||
**问题：** vLLM 遇到不兼容的 tokenizer class（如 `TokenizersBackend`、`TiktokenTokenizer`）时无法加载模型。
|
||
|
||
**解决方案：** 容器启动时自动检测并按需修复 tokenizer 配置：
|
||
|
||
- **仅以下情况才会覆盖** `tokenizer_class`：
|
||
- 值在已知坏列表中（`TokenizersBackend`、`TiktokenTokenizer`）
|
||
- 值缺失或为空
|
||
- 覆盖时根据模型目录中的文件类型推断兜底值：
|
||
|
||
| 文件判断依据 | 替换为 |
|
||
|---|---|
|
||
| 存在 `tokenizer.json` | `PreTrainedTokenizerFast` |
|
||
| 存在 `tokenizer.model` | `LlamaTokenizer` |
|
||
| 存在 `vocab.json` + `merges.txt` | `GPT2TokenizerFast` |
|
||
| 无法识别 | `PreTrainedTokenizerFast` |
|
||
|
||
**执行流程：**
|
||
|
||
```
|
||
容器启动 → entrypoint.sh → fix_tokenizer.py（按需修复）→ vllm serve --tokenizer /tmp/fixed_tokenizer
|
||
```
|
||
|
||
**已验证模型：**
|
||
- ChuGyouk/F_R9_T3 ✅（原 `TokenizersBackend` 报错已修复）
|
||
|
||
### Patch 3：非标准 head_size 自动切换 Attention Backend
|
||
|
||
**问题：** FlashAttention 仅支持白名单内的 head_size `[32, 64, 96, 128, 160, 192, 224, 256]`，head_size=80 或 100 的模型直接报错拒绝启动。
|
||
|
||
**根本原因：** FlashAttention 是针对固定 head_size 的手写 kernel，没有对应 kernel 就无法运行。
|
||
|
||
**解决方案：** 两步组合修复：
|
||
|
||
1. **构建时**（`patch_triton.py`）：将 Triton backend 的 `validate_head_size` 函数改为直接 `return`，解除其白名单限制（Triton 是运行时编译，实际可支持任意 head_size）
|
||
|
||
2. **运行时**（`detect_head_size.py` + `entrypoint.sh`）：读取模型 `config.json` 计算 head_size，若不在白名单则设置 `VLLM_USE_FLASH_ATTN_PA=0`，ROCm V1 engine 自动路由到 Triton backend
|
||
|
||
**head_size 计算逻辑：**
|
||
```
|
||
优先读 config.json["head_dim"]
|
||
没有则计算 hidden_size // num_attention_heads
|
||
```
|
||
|
||
**已验证模型：**
|
||
|
||
| 模型 | head_size | 结果 |
|
||
|------|-----------|------|
|
||
| jondurbin/airoboros-3b-3p11 | 80 | ✅ 启动 + 推理正常 |
|
||
| KnutJaegersberg/black_goo_recipe_d | 100 | ✅ 启动 + 推理正常 |
|
||
|
||
## 构建方法
|
||
|
||
```bash
|
||
docker build -t git.modelhub.org.cn:9443/enginex-hygon/vllm:0.9.2-patch-v2 .
|
||
```
|
||
|
||
## 启动示例
|
||
|
||
```bash
|
||
docker run -dit \
|
||
-p 26381:20644 \
|
||
--privileged --shm-size=64G \
|
||
--device=/dev/kfd --device=/dev/dri/ \
|
||
--cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
|
||
--ulimit memlock=-1:-1 --ipc=host --group-add video \
|
||
-v /opt/hyhal:/opt/hyhal \
|
||
-v /path/to/model:/model \
|
||
-e HIP_VISIBLE_DEVICES=0 \
|
||
git.modelhub.org.cn:9443/enginex-hygon/vllm:0.9.2-patch-v2 \
|
||
--port 20644 --served-model-name llm --max-model-len 4096 \
|
||
--enforce-eager --trust-remote-code -tp 1
|
||
```
|
||
|
||
## 文件说明
|
||
|
||
| 文件 | 执行时机 | 作用 |
|
||
|------|---------|------|
|
||
| `Dockerfile` | — | 镜像构建配置，集成全部修复方案 |
|
||
| `patch.py` | 构建时 | 修复 transformers `extra_special_tokens` list/dict 兼容问题 |
|
||
| `patch_triton.py` | 构建时 | 绕过 Triton backend 的 head_size 白名单验证 |
|
||
| `entrypoint.sh` | 运行时 | 容器入口，协调 tokenizer 修复和 head_size 检测，再启动 vLLM |
|
||
| `detect_tokenizer.py` | 运行时 | 根据模型目录文件推断 tokenizer 类型 |
|
||
| `fix_tokenizer.py` | 运行时 | 将修复后的 tokenizer 配置写入 `/tmp/fixed_tokenizer/` |
|
||
| `detect_head_size.py` | 运行时 | 读取 `config.json` 判断 head_size 是否在 FlashAttention 白名单内 |
|
||
|
||
## 注意事项
|
||
|
||
1. **原始模型文件不受影响**：tokenizer 修复在 `/tmp/fixed_tokenizer/` 中进行，不修改挂载的模型目录。
|
||
|
||
## 参考资料
|
||
|
||
- `NV A100 Patched 镜像合并`
|
||
- 原始 Tokenizer 修复方案：[enginex-s2-sglang-private](https://dev.modelhub.org.cn/EngineX-Sunrise/enginex-s2-sglang-private)
|
||
- vLLM 官方文档：https://docs.vllm.ai/
|
||
|
||
## 维护信息
|
||
|
||
- **构建日期**：2026-05-19
|
||
- **维护者**：i-peixingyu |