commit 60745be0e2db7a3ae212ec23c68cac0a5761fb08 Author: i-peixingyu Date: Tue Jul 14 01:20:25 2026 +0800 上传文件至 / diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..713fda9 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,21 @@ +FROM harbor.4pd.io/dooke/vllm/vllm/vllm-openai:v0.11.0 + +# 修复1 transformers 库级 patch +COPY patch.py /tmp/patch.py +RUN python3 /tmp/patch.py + +# 修复2 Triton backend head_size 验证绕过 +COPY patch_triton.py /tmp/patch_triton.py +RUN python3 /tmp/patch_triton.py + +# 修复3 运行时 tokenizer 配置修复脚本 +COPY detect_tokenizer.py /opt/detect_tokenizer.py +COPY fix_tokenizer.py /opt/fix_tokenizer.py + +# 修复4 head_size 检测,自动切换 attention backend +COPY detect_head_size.py /opt/detect_head_size.py + +COPY entrypoint.sh /opt/entrypoint.sh +RUN chmod +x /opt/entrypoint.sh + +ENTRYPOINT ["/opt/entrypoint.sh"] diff --git a/detect_head_size.py b/detect_head_size.py new file mode 100644 index 0000000..976750c --- /dev/null +++ b/detect_head_size.py @@ -0,0 +1,27 @@ +import json, os, sys + +MODEL_DIR = os.environ.get("MODEL_DIR", "/model") +cfg_path = os.path.join(MODEL_DIR, "config.json") + +if not os.path.exists(cfg_path): + sys.exit(0) + +with open(cfg_path) as f: + cfg = json.load(f) + +head_size = cfg.get("head_dim") +if head_size is None: + hs = cfg.get("hidden_size") + nh = cfg.get("num_attention_heads") + if hs and nh: + head_size = hs // nh + +if head_size is None: + sys.exit(0) + +SUPPORTED = {32, 64, 96, 128, 160, 192, 224, 256} +if head_size not in SUPPORTED: + print(head_size) + sys.exit(2) + +sys.exit(0) diff --git a/detect_tokenizer.py b/detect_tokenizer.py new file mode 100644 index 0000000..03396c2 --- /dev/null +++ b/detect_tokenizer.py @@ -0,0 +1,25 @@ +import os +import json + +def detect(model_dir): + cfg_path = os.path.join(model_dir, "tokenizer_config.json") + + if os.path.exists(cfg_path): + with open(cfg_path) as f: + cfg = json.load(f) + cls = cfg.get("tokenizer_class", "") + else: + cls = "" + + files = os.listdir(model_dir) + + if "tokenizer.json" in files: + return "fast", cls + + if "tokenizer.model" in files: + return "sentencepiece", cls + + if "vocab.json" in files and "merges.txt" in files: + return "bpe", cls + + return "unknown", cls \ No newline at end of file diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100644 index 0000000..090ad34 --- /dev/null +++ b/entrypoint.sh @@ -0,0 +1,24 @@ +#!/bin/bash +set -e + +MODEL_DIR=${MODEL_DIR:-/model} +FIX_TOKENIZER_DIR=/tmp/fixed_tokenizer + +echo "[entrypoint] fixing tokenizer..." +python3 /opt/fix_tokenizer.py + +echo "[entrypoint] checking head_size..." +set +e +HEAD_OUT=$(python3 /opt/detect_head_size.py) +RC=$? +set -e + +if [ "$RC" = "2" ] && [ -n "$HEAD_OUT" ]; then + export VLLM_USE_FLASH_ATTN_PA=0 + echo "[entrypoint] head_size=$HEAD_OUT not in FlashAttention whitelist, switching to Triton backend (VLLM_USE_FLASH_ATTN_PA=0)" +fi + +echo "[entrypoint] starting vllm..." +exec vllm serve "$MODEL_DIR" \ + --tokenizer "$FIX_TOKENIZER_DIR" \ + "$@" diff --git a/fix_tokenizer.py b/fix_tokenizer.py new file mode 100644 index 0000000..2981d52 --- /dev/null +++ b/fix_tokenizer.py @@ -0,0 +1,64 @@ +import os +import shutil +import json +import transformers +import inspect +from detect_tokenizer import detect + +MODEL_DIR = os.environ.get("MODEL_DIR", "/model") +OUT_DIR = os.environ.get("FIX_TOKENIZER_DIR", "/tmp/fixed_tokenizer") + +os.makedirs(OUT_DIR, exist_ok=True) + +def copy_if_exists(name): + src = os.path.join(MODEL_DIR, name) + if os.path.exists(src): + shutil.copy(src, OUT_DIR) + +for f in [ + "tokenizer.json", + "tokenizer_config.json", + "special_tokens_map.json", + "vocab.json", + "merges.txt", + "tokenizer.model", +]: + copy_if_exists(f) + +typ, orig_cls = detect(MODEL_DIR) + +cfg_path = os.path.join(OUT_DIR, "tokenizer_config.json") + +if os.path.exists(cfg_path): + with open(cfg_path) as f: + cfg = json.load(f) +else: + cfg = {} + +VALID_CLASSES = { + name for name, obj in inspect.getmembers(transformers) + if inspect.isclass(obj) and "Tokenizer" in name +} + +BAD_CLASSES = {"TokenizersBackend", "TiktokenTokenizer"} + +FALLBACK = { + "fast": "PreTrainedTokenizerFast", + "sentencepiece": "LlamaTokenizer", + "bpe": "GPT2TokenizerFast", +} + +if orig_cls and orig_cls in VALID_CLASSES and orig_cls not in BAD_CLASSES: + print(f"[fix_tokenizer] tokenizer_class '{orig_cls}' is valid, skip override") +else: + fallback = FALLBACK.get(typ, "PreTrainedTokenizerFast") + if orig_cls: + print(f"[fix] override bad tokenizer_class: {orig_cls} → {fallback}") + else: + print(f"[fix] tokenizer_class missing, set to: {fallback}") + cfg["tokenizer_class"] = fallback + +with open(cfg_path, "w") as f: + json.dump(cfg, f) + +print(f"[fix_tokenizer] done → {OUT_DIR}") \ No newline at end of file