feat: Embedding FP16 开关 + 团队成员信息完善 + gitignore 更新
- infer.py: 新增 emb_fp16 CONFIG 选项(默认 False),Embedding 权重可 FP16 省查表带宽 - CLAUDE.md: 补充团队成员表(AI Studio 用户名→真实姓名) - README.md: 新增团队区块,标注三人参赛身份 - .gitignore: 排除 DVC/HF 工具自动生成的元数据文件 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
+9
-4
@@ -34,6 +34,7 @@ except Exception:
|
||||
# ============================================================
|
||||
CONFIG = {
|
||||
"fp16": True, # True=半精度推理;False=FP32 参考跑(确立 AUC 天花板)
|
||||
"emb_fp16": False, # True=Embedding 也 FP16(省 ~10GB 显存带宽,AUC 可能微降)
|
||||
"keep_fp32_modules": (), # fp16 下仍保留 FP32 的子模块名前缀,如 ("linear",)
|
||||
"expert_merge": True, # 是否做 expert 权重相似度合并
|
||||
"merge_threshold": 0.90, # 合并的余弦相似度阈值
|
||||
@@ -700,14 +701,18 @@ def load_model(ckpt_path, device='cuda:0'):
|
||||
|
||||
if CONFIG["fp16"]:
|
||||
model = model.half()
|
||||
# Embedding 始终保留 FP32(int 索引查表,不受浮点精度影响)
|
||||
model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)
|
||||
# Embedding FP16:省 ~50% 查表带宽(5M×512: 10GB→5GB),AUC 可能微降
|
||||
if not CONFIG.get("emb_fp16", False):
|
||||
model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)
|
||||
# 额外保留 FP32 的精度敏感模块(输入/输出自动转换)
|
||||
for name, module in model.named_modules():
|
||||
if name and any(name.startswith(p) for p in CONFIG["keep_fp32_modules"]):
|
||||
_force_fp32_io(module)
|
||||
print(f"[INFO] FP16 on; FP32-kept: "
|
||||
f"{('rep_encoder.emb',) + tuple(CONFIG['keep_fp32_modules'])}")
|
||||
kept = []
|
||||
if not CONFIG.get("emb_fp16", False):
|
||||
kept.append("rep_encoder.emb")
|
||||
kept.extend(CONFIG["keep_fp32_modules"])
|
||||
print(f"[INFO] FP16 on; FP32-kept: {tuple(kept)}")
|
||||
else:
|
||||
model = model.float()
|
||||
print("[INFO] FP32 reference (no half)")
|
||||
|
||||
Reference in New Issue
Block a user