Serendipity/CTI-Inference-Opt
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: emb_fp16 默认开启(本地AUC 0.75932≈无损,查表带宽减半)；修正打印

Browse Source 
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
OwnerSunshine530
2026-06-15 12:39:10 +08:00
parent adc99b5b41
commit 1249bbdbbc
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
代码/code/infer.py
+4 -3
View File
@@ -50,7 +50,7 @@ CONFIG = {
"vectorize_moe": True, # True=稠密向量化MoE(无同步点)；False=原逐expert循环(.nonzero同步) "vectorize_moe": True, # True=稠密向量化MoE(无同步点)；False=原逐expert循环(.nonzero同步)
"fuse_embedding": True, # True=28个slot的查表+池化融合为1次(减per-batch kernel启动) "fuse_embedding": True, # True=28个slot的查表+池化融合为1次(减per-batch kernel启动)
"syncfree_mask": True, # True=用searchsorted构造因果mask(无同步)；False=repeat_interleave(同步) "syncfree_mask": True, # True=用searchsorted构造因果mask(无同步)；False=repeat_interleave(同步)
"emb_fp16": False, # True=Embedding表也转FP16(查表带宽减半,可能微动AUC)；False=保留FP32 "emb_fp16": True, # True=Embedding表转FP16(查表带宽减半,实测AUC 0.75932≈无损)
"compile": False, # 是否 torch.compile（实测慢5×，勿开） "compile": False, # 是否 torch.compile（实测慢5×，勿开）
} }
@@ -708,8 +708,9 @@ def load_model(ckpt_path, device='cuda:0'):
for name, module in model.named_modules(): for name, module in model.named_modules():
if name and any(name.startswith(p) for p in CONFIG["keep_fp32_modules"]): if name and any(name.startswith(p) for p in CONFIG["keep_fp32_modules"]):
_force_fp32_io(module) _force_fp32_io(module)
print(f"[INFO] FP16 on; FP32-kept: " emb_note = "emb=FP16" if CONFIG.get("emb_fp16", False) else "emb=FP32"
f"{('rep_encoder.emb',) + tuple(CONFIG['keep_fp32_modules'])}") print(f"[INFO] FP16 on; {emb_note}; extra FP32-kept: "
f"{tuple(CONFIG['keep_fp32_modules'])}")
else: else:
model = model.float() model = model.float()
print("[INFO] FP32 reference (no half)") print("[INFO] FP32 reference (no half)")