diff --git a/docs/superpowers/plans/2026-06-14-cti-auc-recovery.md b/docs/superpowers/plans/2026-06-14-cti-auc-recovery.md new file mode 100644 index 0000000..1339983 --- /dev/null +++ b/docs/superpowers/plans/2026-06-14-cti-auc-recovery.md @@ -0,0 +1,821 @@ +# CTI 推理优化冲击 80+ 实现计划 + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** 在不改模型结构、不训练测试集的前提下,先找回当前推理丢失的 AUC,再做结构性延迟重写,把榜上分数从 58.86 推向 80+。 + +**Architecture:** 在 AI Studio notebook(A800 + dataset + ckpt.pt)里,先建一个带同步计时和配置开关的测量闭环 `bench.py`;阶段 A 用消融实验定位并找回 AUC(30 分桶);阶段 B 用数值等价的内核重写压低延迟(块对角注意力 / MoE 向量化 / embedding 融合)。每步过本地关卡,再用有限的提交确认验证集。 + +**Tech Stack:** Python 3.10, PyTorch 2.6.0 (CUDA 12.4), NVIDIA A800 (SM80), sklearn (AUC), AI Studio notebook。 + +--- + +## 执行环境约定 + +- 所有运行都在 **AI Studio notebook** 内(本地 Windows 只装了 numpy+tqdm,跑不了 torch)。 +- 提交文件只有 `infer.py` / `requirements.txt` / `build_env.sh` 会被打包;`bench.py`、`tests/` **绝不进提交包**。 +- 每个改 `infer.py` 的任务,最后都要确认 `bench.py` 默认配置仍能复现「当前最优」,避免污染提交版本。 +- 数据路径(notebook 内):`代码/code/dataset/`(软链)、`代码/code/ckpt.pt`、本地标签 `dataset/label_data.txt`。 + +## 文件结构 + +| 文件 | 职责 | 是否提交 | +|------|------|----------| +| `代码/code/infer.py` | 提交主脚本。引入模块级 `CONFIG` 开关;`load_model`/`RepEncoder`/`SMoE`/注意力按 `CONFIG` 行为,默认值=当前最优 | ✅ | +| `代码/code/bench.py` | 测量闭环。设置 `infer.CONFIG`,跑本地推理,同步计时,打印 AUC/PCOC/延迟/总分;支持配置扫描 | ❌ | +| `代码/code/tests/test_equiv.py` | 阶段 B 重写的数值等价测试(新实现 vs 原实现 allclose) | ❌ | +| `代码/code/EXPERIMENTS.md` | 实验记录表(配置 → AUC/PCOC/延迟/本地分/提交分) | ❌(可入 git,不入提交包) | + +--- + +## 阶段 0:测量闭环 + +### Task 1: 给 infer.py 加 CONFIG 开关板 + +**Files:** +- Modify: `代码/code/infer.py`(顶部新增 CONFIG;改 `load_model`、`RepEncoder.forward`) + +- [ ] **Step 1: 在 import 之后、数据加载层之前插入模块级 CONFIG** + +```python +# ============================================================ +# 实验配置开关(提交时保持默认 = 当前最优行为) +# bench.py 会在 import 后覆盖这些值;评测系统不碰它,用默认值。 +# ============================================================ +CONFIG = { + "fp16": True, # True=半精度;False=FP32 参考 + "keep_fp32_modules": (), # 在 fp16 下仍保留 FP32 的子模块名前缀,如 ("rep_encoder.emb",) + "expert_merge": True, # 是否做 expert 相似度合并 + "merge_threshold": 0.90, # 合并余弦阈值 + "signid_mode": "clamp", # "clamp" 或 "modulo",处理超界 sign id + "sync_timing": False, # bench 里设 True,做 torch.cuda.synchronize 真实计时 +} +``` + +- [ ] **Step 2: 改 `RepEncoder.forward`,按 CONFIG 处理 sign id** + +把 `代码/code/infer.py` 中 `RepEncoder.forward` 的这一行: + +```python + values = values.clamp(0, max_idx) # 超出 vocab_size 的 sign id 截断,避免越界 +``` + +替换为: + +```python + if CONFIG["signid_mode"] == "modulo": + values = values % self.emb.num_embeddings + else: + values = values.clamp(0, max_idx) +``` + +- [ ] **Step 3: 改 `load_model`,按 CONFIG 控制 fp16 / 保留 FP32 模块 / expert 合并** + +把 `load_model` 中从 `model = model.half()` 到 `_merge_experts(...)` 这一段: + +```python + # === FP16 量化:模型参数转半精度,Embedding 保留 FP32 === + model = model.half() + model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32) + print("[INFO] Model converted to FP16 (embedding kept in FP32)") + + # === 按 Expert 权重相似度合并冗余 expert === + _merge_experts(model, sim_threshold=0.90) +``` + +替换为: + +```python + if CONFIG["fp16"]: + model = model.half() + # embedding 始终保留 FP32(int 索引查表) + model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32) + # 额外保留 FP32 的模块(精度敏感层) + for name, module in model.named_modules(): + if any(name.startswith(p) for p in CONFIG["keep_fp32_modules"]): + module.to(torch.float32) + print(f"[INFO] FP16 on; FP32-kept: {('rep_encoder.emb',) + CONFIG['keep_fp32_modules']}") + else: + model = model.float() + print("[INFO] FP32 reference (no half)") + + if CONFIG["expert_merge"]: + _merge_experts(model, sim_threshold=CONFIG["merge_threshold"]) + else: + print("[INFO] expert_merge off") +``` + +注意:`keep_fp32_modules` 里若含某层(如 `seq_encoder.norm1`),其输入需在该层处转回 FP32。先只用整体 fp16/fp32 与 emb,敏感层在 Task 5 单独处理;本任务只接好开关。 + +- [ ] **Step 4: 在 notebook 跑一遍默认配置,确认行为未变** + +Run(notebook cell): +```python +%cd /home/aistudio/code +!python infer.py +``` +Expected:打印 `FP16 on`、expert 合并日志,AUC ≈ 0.759、PCOC ≈ 1.05~1.11(与改动前一致,证明开关默认值没改变行为)。 + +- [ ] **Step 5: Commit** + +```bash +git add 代码/code/infer.py +git commit -m "feat: infer.py 增加 CONFIG 实验开关(默认=当前最优行为)" +``` + +### Task 2: 建 bench.py 测量闭环 + +**Files:** +- Create: `代码/code/bench.py` + +- [ ] **Step 1: 写 bench.py** + +```python +"""本地测量闭环:设置 infer.CONFIG,跑推理,同步计时,打印指标。不进提交包。""" +import sys, time, io +from pathlib import Path +import torch +from torch.utils.data import DataLoader + +import infer # 同目录 + + +def run_once(config_override: dict, batch_size: int = 50, max_batches: int | None = None): + infer.CONFIG.update(config_override) + infer.CONFIG["sync_timing"] = True + + cur = Path(__file__).parent + ref = cur / "dataset" + history = ref / "history" + test_csv = ref / "test.csv" + label_file = ref / "label_data.txt" + + files = (sorted(history.glob("*.csv")) if history.exists() else []) + [test_csv] + item_dict, user_seq = infer.load_sample_files(files) + test_logids = infer.load_logids_from_file(test_csv) + ds = infer.CTRTestSeqDataset( + test_logids_ordered=list(test_logids), item_dict=item_dict, + user_seq=user_seq, max_feasign_per_slot={1: 2}, max_ctx_len=None, + ) + loader = DataLoader(ds, batch_size=batch_size, shuffle=False, num_workers=0, + collate_fn=infer.make_collate_fn(ds.max_slot_id)) + batches = [] + for b in loader: + batches.append(infer.move_batch_to_device(b, torch.device("cpu"))) + if max_batches and len(batches) >= max_batches: + break + + model, dev = infer.load_model(ckpt_path=None) + logid2p, t_sum = {}, 0.0 + with torch.inference_mode(): + for b in batches: + b = infer.move_batch_to_device(b, dev) + pm = b["pred_mask"].bool() + torch.cuda.synchronize() + t0 = time.time() + logits, _ = model(b) + probs = torch.sigmoid(logits.squeeze(-1)) + torch.cuda.synchronize() + t_sum += time.time() - t0 + for lid, p in zip(b["logid"][pm].cpu().tolist(), probs[pm].cpu().tolist()): + logid2p[lid] = p + + # 按 test.csv 顺序写 predict 并打分 + order = [int(l.split(",")[0]) for l in open(test_csv) if l.strip()] + pred_path = cur / "predict.txt" + with open(pred_path, "w") as f: + for lid in order: + f.write(f"{logid2p[lid]}\n") + res = infer._cal_score(pred_path, label_file, default_latency=t_sum) + print(f"[BENCH] cfg={config_override} bs={batch_size} -> " + f"AUC={res['auc']:.5f} PCOC={res['pcoc']:.4f} " + f"lat={res['latency']:.2f}s score={res['score_all']:.2f}") + return res + + +if __name__ == "__main__": + run_once({}) # 默认配置基准 +``` + +- [ ] **Step 2: 跑默认配置,建立本地基准** + +Run: +```python +%cd /home/aistudio/code +!python bench.py +``` +Expected:打印 `[BENCH]` 一行,记录 AUC/PCOC/同步后真实延迟/本地分。这是后续所有对比的锚点。 + +- [ ] **Step 3: 建实验记录表并记录第一行** + +Create `代码/code/EXPERIMENTS.md`,写入表头与默认配置那一行(数值用 Step 2 实测填): +```markdown +| 配置 | AUC | PCOC | 延迟(同步) | 本地分 | 提交分 | +|------|-----|------|-----------|--------|--------| +| 默认(当前最优) | <实测> | <实测> | <实测> | <实测> | 58.86 | +``` + +- [ ] **Step 4: Commit** + +```bash +git add 代码/code/bench.py 代码/code/EXPERIMENTS.md +git commit -m "feat: 新增 bench.py 测量闭环 + 实验记录表" +``` + +--- + +## 阶段 A:找回 AUC(30 分桶,最高优先) + +### Task 3: FP32 参考跑 —— 确立 AUC 天花板(核心前提验证) + +**Files:** +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 跑纯 FP32、不合并 expert、clamp** + +Run(notebook): +```python +import bench +bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "clamp"}) +``` +Expected:打印一行 AUC/PCOC/延迟。**记录这个 AUC** —— 它是当前代码路径下模型的真实可达上限。 + +- [ ] **Step 2: 判定核心前提** + +把结果记入 EXPERIMENTS.md。判定: +- 若 FP32 AUC 明显 > 默认配置 AUC(如 ≥ +0.01)→ 说明 fp16/合并在掉精度,Task 4/5 有收益。 +- 若 FP32 AUC 仍 ≈ 0.759(验证集对应 ~0.7526)→ **当前数据路径触不到更高 AUC**;缺口可能在 sign-id/特征/上下文(Task 3.5/6),或「80 目标」前提存疑,需暂停并与队友/官方答疑核对(见 spec §10)。 + +- [ ] **Step 3: Commit** + +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: FP32 参考跑,记录 AUC 天花板" +``` + +### Task 4: Sign-ID 取模 vs clamp + +**Files:** +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 先查 max_sign_id 是否超 5M 词表** + +Run(notebook): +```python +import infer +from pathlib import Path +files = sorted(Path("dataset/history").glob("*.csv")) + [Path("dataset/test.csv")] +item_dict, user_seq = infer.load_sample_files(files) +mx = max(int(s) for r in item_dict.values() for s in r["signs"].tolist()) +print("max_sign_id =", mx, "vocab =", 5000000, "超界比例可观?", mx >= 5000000) +``` +Expected:打印最大 sign id。若 `mx >= 5_000_000`,clamp 会把大量 id 压到同一行 —— 头号嫌疑成立。 + +- [ ] **Step 2: FP32 下对比 clamp vs modulo** + +Run: +```python +import bench +bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "clamp"}) +bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "modulo"}) +``` +Expected:两行 AUC。 + +- [ ] **Step 3: 判定 + 记录** + +- modulo 的 AUC 明显更高 → 训练用的就是取模哈希,**保留 modulo**(合规:只是正确还原模型输入,不改结构/权重)。 +- 两者相近或 modulo 更差 → 训练用 clamp/或 id 不超界,保留 clamp。 +记入 EXPERIMENTS.md。 + +- [ ] **Step 4: Commit** + +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: sign-id clamp vs modulo 对比" +``` + +### Task 5: 精度摆放(混合精度找回 AUC) + +**Files:** +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 逐步把敏感层保留 FP32,对比 AUC** + +用上一步定下的 `signid_mode`(记为 `SM`),依次跑: +```python +import bench +bench.run_once({"fp16": True, "expert_merge": False, "signid_mode": SM, + "keep_fp32_modules": ()}) # 纯 fp16 +bench.run_once({"fp16": True, "expert_merge": False, "signid_mode": SM, + "keep_fp32_modules": ("linear",)}) # 保留最终输出头 +bench.run_once({"fp16": True, "expert_merge": False, "signid_mode": SM, + "keep_fp32_modules": ("linear", "rep_encoder.input_norm", + "rep_encoder.linear")}) # +RepEncoder 头 +``` +Expected:三行 AUC + 延迟。 + +- [ ] **Step 2: 选「AUC 最接近 FP32 且延迟可接受」的组合** + +记 `KEEP` = 选中的 `keep_fp32_modules`。判定标准:相对 FP32 参考,AUC 损失 ≤ 0.001 优先;若纯 fp16 已无损,则 `KEEP=()`。记入 EXPERIMENTS.md。 + +- [ ] **Step 3: Commit** + +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: 混合精度摆放,确定 keep_fp32_modules" +``` + +### Task 6: Expert 合并的 AUC 代价 + +**Files:** +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 在选定精度下对比 expert_merge 开/关** + +```python +import bench +bench.run_once({"fp16": True, "signid_mode": SM, "keep_fp32_modules": KEEP, + "expert_merge": False}) +bench.run_once({"fp16": True, "signid_mode": SM, "keep_fp32_modules": KEEP, + "expert_merge": True, "merge_threshold": 0.90}) +``` +Expected:两行,含 AUC 与延迟。 + +- [ ] **Step 2: 判定** + +- 合并掉 AUC(> 0.0005)但只省一点延迟 → **关掉合并**(延迟从阶段 B 补,那里不损精度)。 +- 合并不掉 AUC → 保留。记 `MERGE` = 最终决定。记入 EXPERIMENTS.md。 + +- [ ] **Step 3: Commit** + +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: 量化 expert 合并的 AUC 代价并决定开关" +``` + +### Task 7: 特征与上下文完整性核查 + +**Files:** +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 核查 max_feasign_per_slot 截断的影响** + +```python +import bench +bench.run_once({"fp16": True, "signid_mode": SM, "keep_fp32_modules": KEEP, + "expert_merge": MERGE}) # 当前 dataset 用 {1:2} +``` +然后改 bench.run_once 里 `max_feasign_per_slot={1: 2}` 为 `None`(临时编辑 bench.py 或加参数),再跑一次,对比 AUC。 +Expected:两行。若去掉截断 AUC 升高,说明截断在丢信息。 + +> 注意:评测系统构造 `CTRTestSeqDataset` 时传哪些 `max_feasign_per_slot`/`max_ctx_len` 由评测端决定,**我们不一定能控制**。本步先确认「完整特征是否更好」,若是,则在 `CTRTestSeqDataset.__init__` 里对截断做更保守的默认(仅在确证合规、不属"序列截断"违规的前提下)。 + +- [ ] **Step 2: 核查每条测试样本是否 attend 到完整用户历史** + +```python +import infer +from pathlib import Path +files = sorted(Path("dataset/history").glob("*.csv")) + [Path("dataset/test.csv")] +item_dict, user_seq = infer.load_sample_files(files) +test_uids = {item_dict[l]["userid"] for l in infer.load_logids_from_file(Path("dataset/test.csv"))} +have_hist = sum(1 for u in test_uids if len(user_seq.get(u, [])) > 1) +print(f"测试用户 {len(test_uids)},其中有历史序列(>1)的 {have_hist} " + f"({have_hist/len(test_uids):.1%});序列长度分布:") +import numpy as np +lens = np.array([len(user_seq.get(u, [])) for u in test_uids]) +print("min/median/max =", lens.min(), int(np.median(lens)), lens.max()) +``` +Expected:绝大多数测试用户应有较长历史序列。若大量用户只有长度 1(无历史),说明历史没正确挂上 —— 这会严重压低生成式模型 AUC,需排查 `load_sample_files` 的 userid 关联与排序。 + +- [ ] **Step 3: 记录结论 + Commit** + +把两步结论记入 EXPERIMENTS.md。 +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: 特征截断与上下文完整性核查" +``` + +### Task 8: 锁定阶段 A 最优配置并设为 infer.py 默认 + 提交验证 + +**Files:** +- Modify: `代码/code/infer.py`(把 CONFIG 默认值改为阶段 A 选定组合) + +- [ ] **Step 1: 更新 infer.py 的 CONFIG 默认值** + +把 `CONFIG` 默认值改成 Task 4~7 选定的 `signid_mode=SM`、`keep_fp32_modules=KEEP`、`expert_merge=MERGE`、`merge_threshold` 等(`sync_timing` 保持 False)。 + +- [ ] **Step 2: 跑默认配置确认达到阶段 A 最优本地分** + +```python +%cd /home/aistudio/code +!python bench.py +``` +Expected:AUC ≥ 默认基准,本地分高于先前。 + +- [ ] **Step 3: 打包并提交一次(消耗 1 次/天额度)** + +```bash +cd /home/aistudio/code +rm -f predict.txt +zip -y ../eval.zip infer.py requirements.txt build_env.sh +# 确认包内无 dataset/、无 ckpt.pt、无 bench.py/tests/ +unzip -l ../eval.zip +``` +然后在 AI Studio 提交页提交 `eval.zip`。 + +- [ ] **Step 4: 记录验证集分数 + Commit** + +把提交得到的验证集 AUC/PCOC/延迟/分数记入 EXPERIMENTS.md。 +```bash +git add 代码/code/infer.py 代码/code/EXPERIMENTS.md +git commit -m "feat: 锁定阶段A最优配置为默认 + 验证集提交结果" +``` + +--- + +## 阶段 B:结构性延迟重写(数值等价,不动 AUC) + +> 每个重写任务都先写「新实现 vs 原实现 allclose」等价测试,再替换,最后用 bench 确认 AUC 不变、延迟下降。 + +### Task 9: 块对角因果注意力(FlexAttention) + +**Files:** +- Create: `代码/code/tests/test_equiv.py` +- Modify: `代码/code/infer.py`(`scaled_dot_product` / `CTRModel.forward` mask 路径) + +- [ ] **Step 1: 写等价测试(先失败)** + +Create `代码/code/tests/test_equiv.py`: +```python +import torch, torch.nn.functional as F +import sys; sys.path.insert(0, "..") +import infer + +def _dense_attn(q, k, v, mask): + return F.scaled_dot_product_attention(q, k, v, attn_mask=mask.to(q.dtype).bool()) + +def test_flex_matches_dense(): + torch.manual_seed(0) + B, H, S, Dh = 1, 8, 37, 64 + q, k, v = [torch.randn(B, H, S, Dh, device="cuda") for _ in range(3)] + # 构造 3 个用户的 user_offsets:长度 10/12/15 + offsets = torch.tensor([0, 10, 22, 37], device="cuda") + m = infer.CTRModel.get_sequence_causal_mask.__get__(object())(offsets) # 见下 + dense = _dense_attn(q, k, v, m.unsqueeze(0).unsqueeze(0)) + flex = infer.flex_block_causal_attn(q, k, v, offsets) + assert torch.allclose(dense, flex, atol=1e-3, rtol=1e-3), (dense - flex).abs().max() +``` +> 说明:`get_sequence_causal_mask` 是实例方法,测试里改成直接调用一个等价的独立函数 `infer._build_dense_causal_mask(offsets)`(Step 3 会把现有逻辑抽成模块级函数,便于测试与复用)。把上面 `m = ...` 那行改为 `m = infer._build_dense_causal_mask(offsets)`。 + +- [ ] **Step 2: 跑测试确认失败** + +Run: +```python +%cd /home/aistudio/code/tests +!python -m pytest test_equiv.py::test_flex_matches_dense -v +``` +Expected:FAIL(`infer.flex_block_causal_attn` / `_build_dense_causal_mask` 未定义)。 + +- [ ] **Step 3: 在 infer.py 实现 FlexAttention 路径** + +把 `CTRModel.get_sequence_causal_mask` 的逻辑抽为模块级函数,并新增 flex 实现: +```python +from torch.nn.attention.flex_attention import flex_attention, create_block_mask + +def _build_dense_causal_mask(user_offsets): + lengths = user_offsets[1:] - user_offsets[:-1] + idx = torch.repeat_interleave( + torch.arange(lengths.numel(), device=user_offsets.device), lengths) + same = idx.view(1, -1) == idx.view(-1, 1) + causal = torch.tril(torch.ones_like(same, dtype=torch.bool)) + return same & causal + +def flex_block_causal_attn(q, k, v, user_offsets): + S = q.size(-2) + lengths = user_offsets[1:] - user_offsets[:-1] + doc_id = torch.repeat_interleave( + torch.arange(lengths.numel(), device=q.device), lengths) + def mask_mod(b, h, qi, ki): + return (qi >= ki) & (doc_id[qi] == doc_id[ki]) + block_mask = create_block_mask(mask_mod, B=None, H=None, Q_LEN=S, KV_LEN=S, device=q.device) + return flex_attention(q, k, v, block_mask=block_mask) +``` +然后改 `CTRModel.forward`:mask 不再现造稠密矩阵传给 SDPA,而是把 `user_offsets` 透传,调用 `flex_block_causal_attn`。把 `scaled_dot_product` 改为接收 `extension={"user_offsets": ...}` 并走 flex;`get_sequence_causal_mask` 保留供测试/回退。 + +> 兼容性:FlexAttention 要求 q/k/v 为 `[B,H,S,Dh]`(现有 forward 已是该布局)。FP16 下 atol 放宽到 2e-2 重测。 + +- [ ] **Step 4: 跑测试确认通过** + +Run: +```python +!python -m pytest test_equiv.py::test_flex_matches_dense -v +``` +Expected:PASS。 + +- [ ] **Step 5: bench 确认 AUC 不变、延迟下降** + +```python +import bench, importlib, infer; importlib.reload(infer); importlib.reload(bench) +bench.run_once({}) +``` +Expected:AUC 与 Task 8 一致(±0.0005),延迟较 Task 8 下降。记入 EXPERIMENTS.md。 + +- [ ] **Step 6: Commit** + +```bash +git add 代码/code/infer.py 代码/code/tests/test_equiv.py 代码/code/EXPERIMENTS.md +git commit -m "perf: 块对角因果注意力改用 FlexAttention(数值等价,提速)" +``` + +### Task 10: MoE 向量化(消除 Python 循环与同步) + +**Files:** +- Modify: `代码/code/infer.py`(`SMoE.__init__` 预堆叠权重;`SMoE.forward` 稠密批量计算) +- Modify: `代码/code/tests/test_equiv.py`(加 MoE 等价测试) + +- [ ] **Step 1: 写 MoE 等价测试(先失败)** + +在 `test_equiv.py` 追加: +```python +def test_smoe_vectorized_matches_loop(): + torch.manual_seed(0) + m = infer.SMoE(d_model=512, dim_ff=1024, num_experts=8, k=2).cuda().eval() + x = torch.randn(1, 50, 512, device="cuda") + with torch.no_grad(): + ref, _ = infer._smoe_forward_loop(m, x) # 原实现(保留为参考函数) + new, _ = m(x) # 新向量化实现 + assert torch.allclose(ref, new, atol=1e-4, rtol=1e-4), (ref - new).abs().max() +``` + +- [ ] **Step 2: 跑测试确认失败** + +Run:`!python -m pytest test_equiv.py::test_smoe_vectorized_matches_loop -v` +Expected:FAIL(`_smoe_forward_loop` 未定义 / 新旧不一致)。 + +- [ ] **Step 3: 实现向量化 SMoE** + +把现有 `SMoE.forward` 的循环体抽成模块级 `_smoe_forward_loop(moe, x)`(保留作参考/回退),新 `forward` 改为稠密批量(8 个小 FFN 全算,再按 top-k 选取加权 —— 数学等价,GPU 上无 gather/同步更快): +```python +class SMoE(nn.Module): + def __init__(self, d_model, dim_ff, num_experts, k=2): + super().__init__() + self.num_experts = num_experts + self.k = k + self.experts = nn.ModuleList([Expert(d_model, dim_ff) for _ in range(num_experts)]) + self.gate = TopKGate(d_model, num_experts, k=k) + self._stacked = False + + def _stack_weights(self): + self.register_buffer("W1", torch.stack([e.fc1.weight for e in self.experts])) # [E,F,D] + self.register_buffer("b1", torch.stack([e.fc1.bias for e in self.experts])) # [E,F] + self.register_buffer("W2", torch.stack([e.fc2.weight for e in self.experts])) # [E,D,F] + self.register_buffer("b2", torch.stack([e.fc2.bias for e in self.experts])) # [E,D] + self._stacked = True + + def forward(self, x): + if not self._stacked: + self._stack_weights() + B, S, D = x.shape + topk_idx, topk_score, probs = self.gate(x) + xf = x.reshape(-1, D) # [N,D] + h = torch.einsum("nd,efd->enf", xf, self.W1) + self.b1[:, None, :] # [E,N,F] + h = F.relu(h) + o = torch.einsum("enf,eDf->enD", h, self.W2) + self.b2[:, None, :] # [E,N,D] + o = o.permute(1, 0, 2) # [N,E,D] + idx = topk_idx.reshape(-1, self.k) # [N,k] + sc = topk_score.reshape(-1, self.k) # [N,k] + sel = torch.gather(o, 1, idx.unsqueeze(-1).expand(-1, -1, D)) # [N,k,D] + out = (sel * sc.unsqueeze(-1)).sum(1).reshape(B, S, D) + moe_loss = probs.sum(dim=(0, 1)).std() / (probs.sum(dim=(0, 1)).mean() + 1e-6) + return out, moe_loss +``` +> 注意:合并 expert(Task 6 若开启)会改变 `num_experts` 和权重 —— `_stack_weights` 必须在合并之后、首次 forward 时调用(上面 lazy 实现已满足)。dtype 要与 x 一致(fp16 时 stack 出来即 fp16)。 + +- [ ] **Step 4: 跑测试确认通过** + +Run:`!python -m pytest test_equiv.py::test_smoe_vectorized_matches_loop -v` +Expected:PASS。 + +- [ ] **Step 5: bench 确认 AUC 不变、延迟下降** + +```python +import bench, importlib, infer; importlib.reload(infer); importlib.reload(bench) +bench.run_once({}) +``` +Expected:AUC 一致,延迟较 Task 9 下降。记入 EXPERIMENTS.md。 + +- [ ] **Step 6: Commit** + +```bash +git add 代码/code/infer.py 代码/code/tests/test_equiv.py 代码/code/EXPERIMENTS.md +git commit -m "perf: SMoE 稠密向量化(数值等价,消除循环/同步)" +``` + +### Task 11: Embedding 池化融合(28 次 segment_reduce → 1 次) + +**Files:** +- Modify: `代码/code/infer.py`(`RepEncoder.forward`) +- Modify: `代码/code/tests/test_equiv.py` + +- [ ] **Step 1: 写等价测试(先失败)** + +在 `test_equiv.py` 追加,对比融合实现与逐 slot 实现在同一输入上的输出 allclose(构造一个 28-slot 的小 batch dict,调用 `infer._rep_forward_perslot(enc, batch)` 参考实现 vs `enc(batch)`)。 +```python +def test_rep_fused_matches_perslot(): + torch.manual_seed(0) + enc = infer.RepEncoder(vocab_size=1000, emb_dim=512, slot_num=28, d_model=512).cuda().eval() + batch = {} + for s in range(1, 29): + n = torch.randint(1, 5, (10,)) # 每样本 1~4 个 sign + vals = torch.randint(0, 1000, (int(n.sum()),)) + offs = torch.cat([torch.zeros(1, dtype=torch.long), n.cumsum(0)]) + batch[s] = (vals.cuda(), offs.cuda()) + with torch.no_grad(): + ref = infer._rep_forward_perslot(enc, batch) + new = enc(batch) + assert torch.allclose(ref, new, atol=1e-4), (ref - new).abs().max() +``` + +- [ ] **Step 2: 跑测试确认失败** + +Run:`!python -m pytest test_equiv.py::test_rep_fused_matches_perslot -v` +Expected:FAIL(`_rep_forward_perslot` 未定义)。 + +- [ ] **Step 3: 实现融合** + +把现有逐 slot 循环抽为 `_rep_forward_perslot(enc, batch)`(参考/回退)。新 `RepEncoder.forward` 把 28 个 slot 的 `values` 拼成一条,offsets 平移拼接成覆盖 `28*N` 段的单一 offsets,一次 `segment_reduce`,再 reshape `[28, N, emb]` → permute/cat 成 `[N, 28*emb]`: +```python +def forward(self, batch): + max_idx = self.emb.num_embeddings - 1 + target_dtype = self.input_norm.weight.dtype + N = batch[1][1].numel() - 1 # 样本数 = offsets 段数 + all_vals, seg_offsets, base = [], [0], 0 + for s in range(1, self.slot_num + 1): + vals, offs = batch[s] + if CONFIG["signid_mode"] == "modulo": + vals = vals % self.emb.num_embeddings + else: + vals = vals.clamp(0, max_idx) + all_vals.append(vals) + seg_offsets.extend((offs[1:] + base).tolist()) + base += vals.numel() + cat_vals = torch.cat(all_vals) + seg = torch.tensor(seg_offsets, device=cat_vals.device, dtype=torch.long) + emb = self.emb(cat_vals).to(target_dtype) + pooled = torch.segment_reduce(emb, reduce="sum", offsets=seg, initial=0) # [28*N, emb] + pooled = pooled.view(self.slot_num, N, self.emb_dim).permute(1, 0, 2).reshape(N, -1) + return self.linear(self.input_norm(pooled)) +``` +> 验证点:`seg_offsets` 构造正确性强依赖每个 slot 的 offsets 含开头的 0 —— 测试里务必覆盖「某样本某 slot 为空」的情况(offsets 出现连续相等)。FP16 下放宽 atol。 + +- [ ] **Step 4: 跑测试确认通过** + +Run:`!python -m pytest test_equiv.py::test_rep_fused_matches_perslot -v` +Expected:PASS。 + +- [ ] **Step 5: bench 确认 AUC 不变、延迟下降 + Commit** + +```python +import bench, importlib, infer; importlib.reload(infer); importlib.reload(bench) +bench.run_once({}) +``` +Expected:AUC 一致,延迟下降。记入 EXPERIMENTS.md。 +```bash +git add 代码/code/infer.py 代码/code/tests/test_equiv.py 代码/code/EXPERIMENTS.md +git commit -m "perf: RepEncoder 融合 28 次 segment_reduce 为单次" +``` + +### Task 12: 确认 batch_size 控制权并(若可)扫描最优 + +**Files:** +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 判断评测端是否固定 batch_size** + +查 `代码/任务提交接口说明.md` 与 baseline notebook:评测端自建 DataLoader 时 `batch_size` 是否由其设定。若由评测端固定 → 我们无法在评测改 batch(**跳过本任务**,只在本地扫描了解趋势)。若 infer.py 的 `main()` 才建 loader 而评测复用我们的某入口 → 记录可控。 + +- [ ] **Step 2: 本地扫描 batch_size 的延迟趋势** + +```python +import bench +for bs in [50, 100, 200, 400]: + bench.run_once({}, batch_size=bs) +``` +Expected:延迟随 bs 变化曲线(注意显存)。记入 EXPERIMENTS.md,作为「若可控则用」的参考。 + +- [ ] **Step 3: Commit** + +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: batch_size 控制权确认与延迟扫描" +``` + +### Task 13: 重估 torch.compile / CUDA Graph(图理干净后) + +**Files:** +- Modify: `代码/code/infer.py`、`代码/code/build_env.sh` +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 对干净后的模型试 torch.compile** + +在 `load_model` 末尾(`model.eval()` 后)加可开关的: +```python +if CONFIG.get("compile", False): + model = torch.compile(model, mode="max-autotune", dynamic=True) +``` +`build_env.sh` 加预热(按 spec §11 模板)。bench 对比开/关。 +> FlexAttention 与 torch.compile 通常配合良好(flex 本就鼓励 compile);这次重估可能与上次(失败)结果不同。 + +- [ ] **Step 2: bench 对比 + 判定** + +```python +import bench +bench.run_once({"compile": False}) +bench.run_once({"compile": True}) +``` +若 compile 提速且 AUC 不变 → 保留并把 `compile` 默认设 True;否则关掉。CUDA Graph 仅在序列长度分桶后另行评估,本任务不强求。记入 EXPERIMENTS.md。 + +- [ ] **Step 3: Commit** + +```bash +git add 代码/code/infer.py 代码/code/build_env.sh 代码/code/EXPERIMENTS.md +git commit -m "exp: 图清理后重估 torch.compile" +``` + +--- + +## 阶段 C:收尾 + +### Task 14: PCOC 校准(可选,免费零头) + +**Files:** +- Modify: `代码/code/infer.py`(输出处单调缩放) +- Modify: `代码/code/EXPERIMENTS.md` + +- [ ] **Step 1: 在历史数据上估校准系数** + +用带标签的历史数据估一个对 logit 的温度/偏移 `(a, b)`,使 `mean(sigmoid(a*logit+b)) ≈ mean(label)`(只在历史上拟合,**不碰测试集**)。把系数写入 CONFIG(如 `"calib": (a, b)`),在 `CTRModel.forward` 输出前应用:`pred_logits = a * pred_logits + b`(单调,不改 AUC)。 + +- [ ] **Step 2: bench 确认 PCOC 趋近 1、AUC 不变** + +```python +import bench +bench.run_once({}) +``` +Expected:PCOC 更接近 1.0,AUC 不变。记入 EXPERIMENTS.md。 + +- [ ] **Step 3: Commit** + +```bash +git add 代码/code/infer.py 代码/code/EXPERIMENTS.md +git commit -m "feat: 历史数据 PCOC 单调校准(不改 AUC)" +``` + +### Task 15: 最终提交 + 保底 + +**Files:** +- 无代码改动(打包提交) + +- [ ] **Step 1: 全测试 + bench 总确认** + +```python +%cd /home/aistudio/code/tests +!python -m pytest -v +%cd /home/aistudio/code +!python bench.py +``` +Expected:所有等价测试 PASS;本地分为历史最高。 + +- [ ] **Step 2: 打包并校验包内容** + +```bash +cd /home/aistudio/code +rm -f predict.txt +zip -y ../eval.zip infer.py requirements.txt build_env.sh +unzip -l ../eval.zip # 确认无 dataset/、ckpt.pt、bench.py、tests/ +``` + +- [ ] **Step 3: 提交并记录;保留保底版本** + +提交 `eval.zip`,把验证集分数记入 EXPERIMENTS.md。若新版翻车,立即回退到已知保底(当前 58.86 对应的 commit)。 +```bash +git add 代码/code/EXPERIMENTS.md +git commit -m "exp: 最终版本提交结果" +git tag best-$(date +%m%d) # 标记当前最优,便于回退 +``` + +--- + +## 自检(计划 vs spec) + +- spec §4 测量闭环 → Task 1–2 ✅ +- spec §5 阶段 A(sign-id/精度/expert合并/特征/上下文)→ Task 3–8 ✅ +- spec §6 阶段 B(注意力/MoE/embedding/batch/compile)→ Task 9–13 ✅ +- spec §7 PCOC 校准 → Task 14 ✅ +- spec §8 合规与提交纪律(10次/天、保底、包校验)→ Task 8/15 ✅ +- spec §9 成功标准(FP32 天花板、≥0.01 AUC 杠杆、延迟≤25s、PCOC∈[0.95,1.05])→ Task 3/4-5/9-13/14 的关卡 ✅ +- spec §10 前提验证(验证集 AUC 是否 > 0.7526)→ Task 3 Step 2 判定门 ✅ + +**已知风险/未决(继承自 spec §10)**: +- 评测端是否固定 `batch_size`、传哪些截断参数 —— Task 7/12 先确认,控制权不在我方则相应任务降级为「仅本地参考」。 +- 核心前提(验证集 AUC 有上行空间)若被 Task 3 证伪,暂停阶段 B,回到与队友/官方答疑核对目标。