i-peixingyu
5c537988f5
All checks were successful
Docker Build and Push / docker (push) Successful in 16m2s
149 lines
5.0 KiB
Markdown
149 lines
5.0 KiB
Markdown
# NV A100 vLLM Patched v2.0
|
||
|
||
基于 `harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0` 构建的 NVidia A100 GPU 专用 vLLM Docker 镜像，包含 4 类兼容性补丁。
|
||
|
||
## 基础镜像
|
||
|
||
- **镜像**: `harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0`
|
||
- **Python**: 3.12
|
||
- **vLLM**: 0.11.0
|
||
|
||
## 补丁概览
|
||
|
||
此镜像在构建时打了 4 类补丁，启动时再执行运行时修复：
|
||
|
||
| # | 文件 | 类型 | 说明 |
|
||
|---|------|------|------|
|
||
| 1 | `patch.py` | 构建时 | transformers `tokenization_utils_base.py` 库级补丁 |
|
||
| 2 | `patch_triton.py` | 构建时 | Triton Attention backend `validate_head_size` 绕过 |
|
||
| 3 | `detect_tokenizer.py` + `fix_tokenizer.py` | 运行时 | 自动检测并修复 tokenizer 配置 |
|
||
| 4 | `detect_head_size.py` | 运行时 | head_size 检测，自动切换 attention backend |
|
||
|
||
---
|
||
|
||
## 修复 1: transformers 库级补丁 (`patch.py`)
|
||
|
||
**问题**: 部分模型的 `tokenizer_config.json` 中 `extra_special_tokens` 字段是 list 而非 dict，导致 `transformers` 加载 tokenizer 时崩溃。
|
||
|
||
**修复位置**: `/usr/local/lib/python3.12/dist-packages/transformers/tokenization_utils_base.py`
|
||
|
||
**逻辑**: 在 `SPECIAL_TOKENS_ATTRIBUTES` 赋值前插入类型检查——若 `special_tokens` 是 list，则转为 `{t: t}` 的 dict 形式：
|
||
|
||
```python
|
||
# PATCH: some models have extra_special_tokens as list instead of dict
|
||
if isinstance(special_tokens, list):
|
||
special_tokens = {t: t for t in special_tokens}
|
||
```
|
||
|
||
---
|
||
|
||
## 修复 2: Triton Attention head_size 验证绕过 (`patch_triton.py`)
|
||
|
||
**问题**: Triton Attention backend 的 `validate_head_size` 方法强制要求 `head_size >= 32`，某些非标准模型的 head_size 不满足此约束会直接报错。
|
||
|
||
**修复位置**: `/usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/triton_attn.py`
|
||
|
||
**逻辑**: 将整个 `validate_head_size` 方法体替换为直接 `return`，完全绕过验证。Triton 本身在运行时编译，实际支持任意 head_size。
|
||
|
||
---
|
||
|
||
## 修复 3: 运行时 tokenizer 配置修复
|
||
|
||
### `detect_tokenizer.py` — 检测器
|
||
|
||
根据模型目录中的文件自动判断 tokenizer 类型：
|
||
|
||
| 特征文件 | 判定类型 |
|
||
|----------|----------|
|
||
| `tokenizer.json` | `fast` |
|
||
| `tokenizer.model` | `sentencepiece` |
|
||
| `vocab.json` + `merges.txt` | `bpe` |
|
||
| 以上均无 | `unknown` |
|
||
|
||
同时读取 `tokenizer_config.json` 中的 `tokenizer_class` 字段作为原始类名。
|
||
|
||
### `fix_tokenizer.py` — 修复器
|
||
|
||
**运行时机**: 容器启动时（entrypoint 第一步）
|
||
|
||
**逻辑**:
|
||
1. 将模型目录中的 tokenizer 相关文件拷贝到 `/tmp/fixed_tokenizer/`
|
||
2. 调用 `detect_tokenizer` 判断 tokenizer 类型
|
||
3. 检查原 `tokenizer_class` 是否有效：
|
||
- 有效类：`transformers` 库中所有包含 "Tokenizer" 的类
|
||
- 无效类（黑名单）：`TokenizersBackend`、`TiktokenTokenizer`
|
||
- 缺失或无效时，按类型回退：
|
||
|
||
| tokenizer 类型 | fallback 类 |
|
||
|----------------|-------------|
|
||
| `fast` | `PreTrainedTokenizerFast` |
|
||
| `sentencepiece` | `LlamaTokenizer` |
|
||
| `bpe` | `GPT2TokenizerFast` |
|
||
| `unknown` | `PreTrainedTokenizerFast` |
|
||
|
||
4. 将修复后的配置写回 `/tmp/fixed_tokenizer/tokenizer_config.json`
|
||
|
||
**环境变量**:
|
||
- `MODEL_DIR`：模型目录，默认 `/model`
|
||
- `FIX_TOKENIZER_DIR`：修复后 tokenizer 输出目录，默认 `/tmp/fixed_tokenizer`
|
||
|
||
---
|
||
|
||
## 修复 4: head_size 检测 (`detect_head_size.py`)
|
||
|
||
**运行时机**: 容器启动时（entrypoint 第二步）
|
||
|
||
**逻辑**:
|
||
1. 读取 `${MODEL_DIR}/config.json`
|
||
2. 获取 `head_dim`；若无此字段，则用 `hidden_size / num_attention_heads` 计算
|
||
3. 检查是否在 FlashAttention 支持的 head_size 白名单内：
|
||
|
||
```
|
||
{32, 64, 96, 128, 160, 192, 224, 256}
|
||
```
|
||
|
||
4. **退出码**:
|
||
- `0` — head_size 在白名单内，无需处理
|
||
- `2` — head_size 不在白名单，将数值输出到 stdout，entrypoint 会据此切换 backend
|
||
|
||
**entrypoint 联动**: 当 `detect_head_size.py` 返回退出码 2 时，entrypoint 自动设置 `VLLM_USE_FLASH_ATTN_PA=0`，切换到 Triton Attention backend。
|
||
|
||
---
|
||
|
||
## 入口脚本 (`entrypoint.sh`)
|
||
|
||
启动流程：
|
||
|
||
```
|
||
1. python3 /opt/fix_tokenizer.py → 修复 tokenizer 配置
|
||
2. python3 /opt/detect_head_size.py → 检测 head_size 兼容性
|
||
3. 若非白名单 head_size → export VLLM_USE_FLASH_ATTN_PA=0（切到 Triton）
|
||
4. exec vllm serve ${MODEL_DIR} --tokenizer ${FIX_TOKENIZER_DIR} $@
|
||
```
|
||
|
||
启动时自动以修复后的 tokenizer 目录运行 vLLM，用户只需挂载模型到 `/model`。
|
||
|
||
---
|
||
|
||
## 构建
|
||
|
||
```bash
|
||
docker build -t nv-vllm-patched:v2.0 .
|
||
```
|
||
|
||
---
|
||
|
||
## 与 K100 v2.0 的关系
|
||
|
||
此版本是 K100 vLLM Patched v2.0 的 NV A100 移植版。与 K100 版本的主要区别：
|
||
|
||
- 基础镜像相同（`vllm-openai:v0.11.0`）
|
||
- K100 版本面向 K100 芯片集群（含 K100 专属 patch），NV A100 版本去掉了 K100 专属补丁
|
||
- tokenizer 修复和 head_size 检测逻辑保持兼容
|
||
|
||
## 参考
|
||
|
||
- 上游 K100 版本：`k100-vllm-patched-v2.0/K100-vLLM-Patched-v2.0/`
|
||
- 第四范式 modelhub 平台：`modelhub.org.cn`
|
||
- vLLM 官方文档：https://docs.vllm.ai/
|