上传文件至 /

This commit is contained in:
2026-07-14 01:20:25 +08:00
commit 60745be0e2
5 changed files with 161 additions and 0 deletions

21
Dockerfile Normal file
View File

@@ -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"]

27
detect_head_size.py Normal file
View File

@@ -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)

25
detect_tokenizer.py Normal file
View File

@@ -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

24
entrypoint.sh Normal file
View File

@@ -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" \
"$@"

64
fix_tokenizer.py Normal file
View File

@@ -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}")