sunruoxi/enginex-iluvatar-bi100-vllm-fix-tokenizer
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feature:add head size detect and patch some ops
All checks were successful
Docker Build and Push / docker (push) Successful in 1m5s
Details

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-07-27 16:45:05 +08:00
parent e5db2b58e2
commit 0c1bb7a415
6 changed files with 655 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
Dockerfile
View File

@@ -13,5 +13,8 @@ ENV PATH=/root/apps/apache-jmeter-5.6.3/bin:/root/apps/jdk1.8.0_411/bin:/usr/loc
COPY fix_tokenizer.py /opt/
COPY detect_tokenizer.py /opt/
COPY detect_head_size.py /opt/
COPY patch_ops.py /opt/
COPY patched_ops /opt/patched_ops/
COPY entrypoint.sh /opt/
RUN chmod +x /opt/entrypoint.sh

282
README.md
View File

@@ -1,4 +1,6 @@
# vLLM Tokenizer 自动修复方案
# vLLM 自动修复方案
Tokenizer修复 + Head Size补丁 + Ixformer Ops补丁
## 1. 背景
@@ -16,6 +18,10 @@ ValueError: Tokenizer class TokenizersBackend does not exist or is not currently
- 开启 trust_remote_code=True 时，transformers 会强制加载该 class
- vLLM 无法通过参数 override tokenizer class
另外，某些模型的 head_size 可能不在 vLLM 默认支持的列表中（64, 80, 96, 112, 120, 128, 192, 256），导致运行时错误。
此外，某些模型需要特定的 ixformer 操作函数（如 `gelu_tanh_and_mul`），但这些函数可能不在默认的 ixformer 库中。
---
## 2. 方案目标
@@ -27,6 +33,7 @@ ValueError: Tokenizer class TokenizersBackend does not exist or is not currently
无需修改模型文件
无需修改启动命令
自动修复 tokenizer 并启动 vLLM
自动检测并 patch 不支持的 head_size
```
@@ -46,6 +53,12 @@ entrypoint.sh
↓
修复 tokenizer_config.json
↓
检测模型 head_size
↓
如果 head_size 不在支持列表中，patch vLLM 代码
↓
Patch ixformer ops（复制自定义函数并修改 __init__.py）
↓
vllm serve --tokenizer /tmp/fixed_tokenizer
````
@@ -173,12 +186,273 @@ bos_token / eos_token / pad_token 一致
---
## 9. 总结
## 9. Head Size 自动补丁
本方案通过在容器启动阶段引入 tokenizer 修复逻辑，实现：
### 问题
vLLM 默认只支持以下 head_size 值：
```
“模型不动，运行时自适应兼容”
[64, 80, 96, 112, 120, 128, 192, 256]
```
如果模型的 head_size 不在此列表中，会导致运行时错误。
### 检测逻辑
系统会自动从模型的 `config.json` 中检测 head_size，支持以下多种配置格式：
1. **直接读取 `head_dim` 字段**
```json
{
"head_dim": 128
}
```
2. **从 `hidden_size / num_attention_heads` 计算**
```json
{
"hidden_size": 2048,
"num_attention_heads": 16
}
// head_size = 2048 / 16 = 128
```
3. **从 `n_embd / n_head` 计算（GPTJ等模型）**
```json
{
"n_embd": 2048,
"n_head": 16
}
```
4. **直接读取 `d_kv` 字段（T5等模型）**
```json
{
"d_kv": 128
}
```
### 补丁逻辑
如果检测到的 head_size 不在支持列表中，系统会：
1. **备份原文件**
```
/usr/local/corex/lib64/python3/dist-packages/vllam/attention/ops/paged_attn.py.backup
```
2. **修改 get_supported_head_sizes 方法**
```python
@staticmethod
def get_supported_head_sizes() -> List[int]:
return [64, 80, 96, 112, 120, 128, 192, 256, YOUR_NEW_SIZE]
```
3. **保持列表排序**
新的 head_size 会被插入到正确的位置，保持列表升序排列。
### 日志示例
**无需补丁的情况**
```
[entrypoint] checking model head_size...
[detect_head_size] Found head_dim in config: 128
[detect_head_size] Model head_size: 128
[detect_head_size] head_size 128 is already supported by vLLM, skipping patch
```
**需要补丁的情况**
```
[entrypoint] checking model head_size...
[detect_head_size] Calculated from hidden_size(4096) / num_attention_heads(32) = 128
[detect_head_size] Model head_size: 128
[detect_head_size] head_size 160 is NOT in default supported list: [64, 80, 96, 112, 120, 128, 192, 256]
[detect_head_size] Attempting to patch vLLM...
[patch] Backed up original file to /usr/local/.../paged_attn.py.backup
[patch] Successfully added head_size 160 to supported list: [64, 80, 96, 112, 120, 128, 160, 192, 256]
[detect_head_size] Successfully patched vLLM to support head_size 160
```
### 补丁恢复
如需恢复原始文件：
```bash
cp /usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py.backup \
/usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py
```
### 容错机制
**重要**：head_size 检测和 patch 功能具有完整的容错机制：
- ✅ **检测失败不影响启动**：如果无法检测 head_size，vLLM 仍会正常启动
- ✅ **Patch 失败不影响启动**：如果 patch 过程出错，vLLM 仍会正常启动
- ✅ **代码异常不影响启动**：如果检测脚本本身出现异常，vLLM 仍会正常启动
这确保了：
```
本来能跑的模型 → 即使 head_size 检测失败 → 仍然能跑
```
### 日志示例
**检测失败的情况（仍会启动vLLM）**
```
[entrypoint] checking model head_size...
[detect_head_size] Error during head_size detection/patch: config.json not found
[detect_head_size] Continuing with vLLM startup anyway...
[entrypoint] head_size check failed, but continuing with vLLM startup
[entrypoint] starting vLLM...
```
---
## 10. Ixformer Ops 自动补丁
### 问题
某些模型需要特定的 ixformer 操作函数，如 `gelu_tanh_and_mul`，但这些函数可能不在默认的 ixformer 库中。
### 通用解决方案
**重要特性**：系统采用**通用扫描机制**，无需每次修改代码：
1. **自动扫描**：自动扫描 `patched_ops` 目录中的所有 `.py` 文件
2. **批量处理**：批量复制所有文件到目标目录
3. **自动生成import**：为每个文件自动生成对应的 `from .xxx import *` 语句
4. **智能去重**：自动检测已存在的import，避免重复添加
### 使用方法
只需将需要补丁的 ops 文件放入 `patched_ops` 目录即可：
```bash
patched_ops/
├── gelu_tanh_and_mul.py # 第一个ops文件
├── another_op.py # 第二个ops文件
├── third_operation.py # 第三个ops文件
└── ...
```
系统会自动：
- 扫描所有 `.py` 文件（排除 `__init__.py`）
- 复制到 `/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/`
- 在 `__init__.py` 中添加对应的 import 语句
### 补丁逻辑
系统会自动执行以下操作：
1. **扫描源目录**
```bash
源目录: /opt/patched_ops/
自动查找所有 .py 文件（排除 __init__.py）
```
2. **批量复制文件**
```
源文件: /opt/patched_ops/*.py
目标: /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/
```
3. **自动生成并添加 import 语句**
```
目标文件: /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/__init__.py
自动生成: from .gelu_tanh_and_mul import *
from .another_op import *
from .third_operation import *
```
4. **自动备份**
在修改前会自动备份原始的 `__init__.py` 文件为 `__init__.py.backup`
### 日志示例
**成功的批量补丁操作**
```
[entrypoint] patching ixformer ops...
[patch_ops] Starting ixformer ops patch...
[patch_ops] Found 3 ops file(s): gelu_tanh_and_mul.py, another_op.py, third_operation.py
[patch_ops] Backed up /usr/local/.../__init__.py to /usr/local/.../__init__.py.backup
[patch_ops] Copied gelu_tanh_and_mul.py to /usr/local/.../ixformer/functions/
[patch_ops] Copied another_op.py to /usr/local/.../ixformer/functions/
[patch_ops] Copied third_operation.py to /usr/local/.../ixformer/functions/
[patch_ops] Added 3 import statement(s) to /usr/local/.../__init__.py
[patch_ops] Successfully patched ixformer ops
[patch_ops] Patch completed successfully
```
**重复运行（已存在import）**
```
[entrypoint] patching ixformer ops...
[patch_ops] Starting ixformer ops patch...
[patch_ops] Found 3 ops file(s): gelu_tanh_and_mul.py, another_op.py, third_operation.py
[patch_ops] Backup already exists: /usr/local/.../__init__.py.backup
[patch_ops] Copied gelu_tanh_and_mul.py to /usr/local/.../ixformer/functions/
[patch_ops] Copied another_op.py to /usr/local/.../ixformer/functions/
[patch_ops] Copied third_operation.py to /usr/local/.../ixformer/functions/
[patch_ops] All imports already exist in /usr/local/.../__init__.py, skipping modification
[patch_ops] Successfully patched ixformer ops
```
**目标目录不存在的情况**
```
[entrypoint] patching ixformer ops...
[patch_ops] Starting ixformer ops patch...
[patch_ops] Target directory not found: /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/
[patch_ops] Patch failed, but this will not prevent vLLM from starting
[entrypoint] ixformer ops patch failed, but continuing with vllm startup
```
### 补丁恢复
如需恢复原始的 `__init__.py` 文件：
```bash
cp /usr/local/corex/lib64/python3/dist-packages/ixformer/functions/__init__.py.backup \
/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/__init__.py
```
### 容错机制
与 head_size 检测一样，ixformer ops 补丁也具有完整的容错机制：
- ✅ **源目录不存在不影响启动**：如果源目录不存在，vLLM 仍会正常启动
- ✅ **目标目录不存在不影响启动**：如果目标目录不存在，vLLM 仍会正常启动
- ✅ **复制失败不影响启动**：如果复制过程出错，vLLM 仍会正常启动
- ✅ **修改失败不影响启动**：如果修改 __init__.py 失败，vLLM 仍会正常启动
- ✅ **部分失败不影响整体**：即使某个文件处理失败，其他文件仍会继续处理
---
## 11. 总结
本方案通过在容器启动阶段引入多种自动修复和补丁逻辑，实现：
```
"模型不动，运行时自适应兼容"
```
主要功能：
- ✅ 自动修复不兼容的 tokenizer 配置
- ✅ 自动检测并补丁不支持的 head_size
- ✅ 自动补丁 ixformer ops（如 gelu_tanh_and_mul）
- ✅ 无需修改模型文件，无需修改启动命令
- ✅ 完全透明，不影响正常模型部署
- ✅ **完整的容错机制，确保本来能跑的模型不受影响**
### 启动流程总结
```
容器启动
↓
修复 tokenizer（如需要）
↓
检查并 patch head_size（如需要）
↓
检查并 patch ixformer ops（如需要）
↓
启动 vLLM
```
每个步骤都有完整的容错机制，确保即使某个步骤失败，也不会影响后续步骤的执行。

186
detect_head_size.py Executable file
View File

@@ -0,0 +1,186 @@
#!/usr/bin/env python3
"""
检测模型的head_size，如果不在vLLM支持列表中，则patch到vllm代码中
"""
import os
import json
import re
import shutil
from typing import Optional, List
# vLLM默认支持的head_size列表
DEFAULT_SUPPORTED_HEAD_SIZES = [64, 80, 96, 112, 120, 128, 192, 256]
# vLLM需要patch的文件路径
VLLM_PAGED_ATTN_PATH = "/usr/local/corex/lib64/python3/dist-packages/vllm/attention/ops/paged_attn.py"
def get_model_head_size(model_dir: str) -> Optional[int]:
"""
从模型config.json中获取head_size
Args:
model_dir: 模型目录路径
Returns:
head_size，如果无法获取则返回None
"""
config_path = os.path.join(model_dir, "config.json")
if not os.path.exists(config_path):
print(f"[detect_head_size] config.json not found in {model_dir}")
return None
try:
with open(config_path, 'r') as f:
config = json.load(f)
# 尝试多种方式获取head_size
# 1. 直接从head_dim字段获取
if 'head_dim' in config:
head_size = config['head_dim']
print(f"[detect_head_size] Found head_dim in config: {head_size}")
return head_size
# 2. 从hidden_size和num_attention_heads计算
if 'hidden_size' in config and 'num_attention_heads' in config:
hidden_size = config['hidden_size']
num_heads = config['num_attention_heads']
head_size = hidden_size // num_heads
print(f"[detect_head_size] Calculated from hidden_size({hidden_size}) / num_attention_heads({num_heads}) = {head_size}")
return head_size
# 3. 对于GPTJ等模型，可能使用n_embd和n_head
if 'n_embd' in config and 'n_head' in config:
n_embd = config['n_embd']
n_head = config['n_head']
head_size = n_embd // n_head
print(f"[detect_head_size] Calculated from n_embd({n_embd}) / n_head({n_head}) = {head_size}")
return head_size
# 4. 对于T5等模型，可能使用d_kv和num_heads
if 'd_kv' in config:
head_size = config['d_kv']
print(f"[detect_head_size] Found d_kv in config: {head_size}")
return head_size
print(f"[detect_head_size] Cannot determine head_size from config.json")
return None
except Exception as e:
print(f"[detect_head_size] Error reading config.json: {e}")
return None
def patch_vllm_head_sizes(new_head_size: int) -> bool:
"""
将新的head_size patch到vLLM的get_supported_head_sizes方法中
Args:
new_head_size: 需要添加的新head_size
Returns:
是否成功patch
"""
if not os.path.exists(VLLM_PAGED_ATTN_PATH):
print(f"[patch] vLLM file not found: {VLLM_PAGED_ATTN_PATH}")
return False
try:
# 备份原文件
backup_path = VLLM_PAGED_ATTN_PATH + ".backup"
if not os.path.exists(backup_path):
shutil.copy2(VLLM_PAGED_ATTN_PATH, backup_path)
print(f"[patch] Backed up original file to {backup_path}")
# 读取文件内容
with open(VLLM_PAGED_ATTN_PATH, 'r') as f:
content = f.read()
# 找到get_supported_head_sizes方法并patch
pattern = r'(@staticmethod\s+def get_supported_head_sizes\(\)\s*->\s*List\[int\]:\s*return\s*\[)([^\]]+)(\])'
match = re.search(pattern, content, re.MULTILINE | re.DOTALL)
if not match:
print(f"[patch] Could not find get_supported_head_sizes method in {VLLM_PAGED_ATTN_PATH}")
return False
prefix = match.group(1)
current_list = match.group(2)
suffix = match.group(3)
# 解析当前列表
current_sizes = [int(x.strip()) for x in current_list.split(',') if x.strip()]
# 添加新的head_size并排序
if new_head_size not in current_sizes:
current_sizes.append(new_head_size)
current_sizes.sort()
# 生成新的列表字符串
new_list_str = ', '.join(str(size) for size in current_sizes)
# 替换原方法
new_method = f"{prefix}{new_list_str}{suffix}"
new_content = re.sub(pattern, new_method, content, flags=re.MULTILINE | re.DOTALL)
# 写回文件
with open(VLLM_PAGED_ATTN_PATH, 'w') as f:
f.write(new_content)
print(f"[patch] Successfully added head_size {new_head_size} to supported list: {current_sizes}")
return True
else:
print(f"[patch] head_size {new_head_size} already in list, skipping")
return True
except Exception as e:
print(f"[patch] Error patching vLLM file: {e}")
return False
def main():
"""主函数"""
try:
model_dir = os.environ.get("MODEL_DIR", "/model")
print(f"[detect_head_size] Checking model in: {model_dir}")
# 获取head_size
head_size = get_model_head_size(model_dir)
if head_size is None:
print(f"[detect_head_size] Could not determine head_size, skipping")
return
print(f"[detect_head_size] Model head_size: {head_size}")
# 检查是否需要patch
if head_size in DEFAULT_SUPPORTED_HEAD_SIZES:
print(f"[detect_head_size] head_size {head_size} is already supported by vLLM, skipping patch")
return
print(f"[detect_head_size] head_size {head_size} is NOT in default supported list: {DEFAULT_SUPPORTED_HEAD_SIZES}")
print(f"[detect_head_size] Attempting to patch vLLM...")
# 执行patch
success = patch_vllm_head_sizes(head_size)
if success:
print(f"[detect_head_size] Successfully patched vLLM to support head_size {head_size}")
else:
print(f"[detect_head_size] Failed to patch vLLM for head_size {head_size}")
except Exception as e:
print(f"[detect_head_size] Error during head_size detection/patch: {e}")
print(f"[detect_head_size] Continuing with vLLM startup anyway...")
if __name__ == "__main__":
try:
main()
except Exception as e:
print(f"[detect_head_size] Unexpected error: {e}")
print(f"[detect_head_size] This will not prevent vLLM from starting")
import sys
sys.exit(0) # 总是成功退出，不阻止vLLM启动

8
entrypoint.sh
View File

@@ -34,6 +34,14 @@ else
TOKENIZER_ARG=""
fi
# 检查并patch head_size支持（即使失败也不影响启动）
echo "[entrypoint] checking model head_size..."
python3 /opt/detect_head_size.py || echo "[entrypoint] head_size check failed, but continuing with vllm startup"
# Patch ixformer ops（即使失败也不影响启动）
echo "[entrypoint] patching ixformer ops..."
python3 /opt/patch_ops.py || echo "[entrypoint] ixformer ops patch failed, but continuing with vllm startup"
echo "[entrypoint] starting vllm..."
exec vllm serve "$MODEL_DIR" $TOKENIZER_ARG "$@"

151
patch_ops.py Executable file
View File

@@ -0,0 +1,151 @@
#!/usr/bin/env python3
"""
Patch ixformer functions - 自动扫描并复制自定义的ops文件并修改__init__.py
支持批量处理patched_ops目录中的所有.py文件，无需每次修改代码。
只需将新的ops文件放入patched_ops目录即可。
"""
import os
import shutil
from typing import List
# 目标路径
TARGET_DIR = "/usr/local/corex/lib64/python3/dist-packages/ixformer/functions/"
TARGET_INIT_FILE = os.path.join(TARGET_DIR, "__init__.py")
SOURCE_DIR = "/opt/patched_ops/"
def get_ops_files(source_dir: str) -> List[str]:
"""
扫描源目录中的所有.py文件（排除__init__.py）
Args:
source_dir: 源目录路径
Returns:
.py文件名列表
"""
if not os.path.exists(source_dir):
print(f"[patch_ops] Source directory not found: {source_dir}")
return []
ops_files = []
for filename in os.listdir(source_dir):
if filename.endswith('.py') and filename != '__init__.py':
ops_files.append(filename)
return ops_files
def generate_import_statement(filename: str) -> str:
"""
根据文件名生成import语句
Args:
filename: .py文件名（不含扩展名）
Returns:
import语句
"""
module_name = filename.replace('.py', '')
return f"from .{module_name} import *"
def patch_ixformer_ops() -> bool:
"""
Patch ixformer functions - 批量复制ops文件并修改__init__.py
Returns:
是否成功patch
"""
try:
print(f"[patch_ops] Starting ixformer ops patch...")
# 检查目标目录是否存在
if not os.path.exists(TARGET_DIR):
print(f"[patch_ops] Target directory not found: {TARGET_DIR}")
return False
# 扫描源目录中的所有ops文件
ops_files = get_ops_files(SOURCE_DIR)
if not ops_files:
print(f"[patch_ops] No ops files found in {SOURCE_DIR}")
return False
print(f"[patch_ops] Found {len(ops_files)} ops file(s): {', '.join(ops_files)}")
# 备份原__init__.py文件
if os.path.exists(TARGET_INIT_FILE):
backup_file = TARGET_INIT_FILE + ".backup"
if not os.path.exists(backup_file):
shutil.copy2(TARGET_INIT_FILE, backup_file)
print(f"[patch_ops] Backed up {TARGET_INIT_FILE} to {backup_file}")
else:
print(f"[patch_ops] Backup already exists: {backup_file}")
# 读取现有的__init__.py内容
existing_imports = set()
if os.path.exists(TARGET_INIT_FILE):
with open(TARGET_INIT_FILE, 'r') as f:
init_content = f.read()
# 提取现有的from .xxx import *语句
for line in init_content.split('\n'):
line = line.strip()
if line.startswith('from .') and 'import *' in line:
existing_imports.add(line)
# 批量复制文件并收集需要添加的import语句
new_imports = []
for filename in ops_files:
source_file = os.path.join(SOURCE_DIR, filename)
target_file = os.path.join(TARGET_DIR, filename)
# 复制文件
shutil.copy2(source_file, target_file)
print(f"[patch_ops] Copied {filename} to {TARGET_DIR}")
# 生成import语句
import_stmt = generate_import_statement(filename)
if import_stmt not in existing_imports:
new_imports.append(import_stmt)
# 修改__init__.py文件，添加新的import语句
if new_imports:
with open(TARGET_INIT_FILE, 'a') as f:
f.write(f"\n# Auto-patched ixformer ops\n")
for import_stmt in new_imports:
f.write(f"{import_stmt}\n")
print(f"[patch_ops] Added {len(new_imports)} import statement(s) to {TARGET_INIT_FILE}")
else:
print(f"[patch_ops] All imports already exist in {TARGET_INIT_FILE}, skipping modification")
print(f"[patch_ops] Successfully patched ixformer ops")
return True
except Exception as e:
print(f"[patch_ops] Error patching ixformer ops: {e}")
return False
def main():
"""主函数"""
try:
success = patch_ixformer_ops()
if success:
print(f"[patch_ops] Patch completed successfully")
else:
print(f"[patch_ops] Patch failed, but this will not prevent vLLM from starting")
except Exception as e:
print(f"[patch_ops] Unexpected error: {e}")
print(f"[patch_ops] This will not prevent vLLM from starting")
if __name__ == "__main__":
import sys
try:
main()
except Exception as e:
print(f"[patch_ops] Unexpected error: {e}")
print(f"[patch_ops] This will not prevent vLLM from starting")
sys.exit(0) # 总是成功退出，不阻止vLLM启动

29
patched_ops/gelu_tanh_and_mul.py Normal file
View File

@@ -0,0 +1,29 @@
import torch
__all__ = ["gelu_tanh_and_mul"]
def gelu_tanh_and_mul(
input: "torch.Tensor",
output: "torch.Tensor" = None
):
assert isinstance(input, torch.Tensor)
if output is None:
output_shape = list(input.shape)
output_shape[-1] = output_shape[-1] // 2
output = input.new_empty(output_shape)
hidden_size = input.shape[-1] // 2
x = input[..., :hidden_size]
gate = input[..., hidden_size:]
result = x * torch.nn.functional.gelu(
gate,
approximate="tanh"
)
output.copy_(result)
return output