diff --git a/Dockerfile b/Dockerfile index 6d5a744..3acd1e5 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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 diff --git a/README.md b/README.md index f70e799..fb489b7 100644 --- a/README.md +++ b/README.md @@ -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 +``` + +每个步骤都有完整的容错机制，确保即使某个步骤失败，也不会影响后续步骤的执行。 diff --git a/detect_head_size.py b/detect_head_size.py new file mode 100755 index 0000000..44f9406 --- /dev/null +++ b/detect_head_size.py @@ -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启动 diff --git a/entrypoint.sh b/entrypoint.sh index 07308d3..b8ba2d1 100644 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -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 "$@" diff --git a/patch_ops.py b/patch_ops.py new file mode 100755 index 0000000..42ec9b8 --- /dev/null +++ b/patch_ops.py @@ -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启动 diff --git a/patched_ops/gelu_tanh_and_mul.py b/patched_ops/gelu_tanh_and_mul.py new file mode 100644 index 0000000..3b251b8 --- /dev/null +++ b/patched_ops/gelu_tanh_and_mul.py @@ -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