From 2004ad6bb84ab4af4b20be73db8646b5f4c18780 Mon Sep 17 00:00:00 2001 From: OwnerSunshine530 Date: Mon, 15 Jun 2026 17:06:56 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E9=A2=84=E8=AE=A1=E7=AE=97RepEncoder?= =?UTF-8?q?=E7=BC=93=E5=AD=98,model(batch)=E6=8C=89logid=20gather=E8=B7=B3?= =?UTF-8?q?=E8=BF=87embedding=E5=B1=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 不计时的load_model里(或bench从batches)预计算所有item的context-free RepEncoder向量, 排序存(sorted_logids,emb);model(batch)用searchsorted gather、缺失回退现算。逐位等价。 预期 model(batch) 48s->~37s->~70。CONFIG.precompute_rep(eval默认True);bench --precompute-rep。 Co-Authored-By: Claude Opus 4.8 --- 代码/code/bench.py | 24 ++++++++++++++- 代码/code/infer.py | 75 +++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 97 insertions(+), 2 deletions(-) diff --git a/代码/code/bench.py b/代码/code/bench.py index 2d0c056..611d8be 100644 --- a/代码/code/bench.py +++ b/代码/code/bench.py @@ -209,8 +209,11 @@ def run_once(config_override=None, batch_size=50, max_batches=None, if max_feasign_per_slot is None: max_feasign_per_slot = {1: 2} + # 本地用已加载的过滤数据自建 rep 缓存，禁止 load_model 自动加载全量数据集 + want_precompute = bool(config_override.pop("precompute_rep", False)) infer.CONFIG.update(config_override) infer.CONFIG["sync_timing"] = True + infer.CONFIG["precompute_rep"] = False cur = Path(__file__).parent ref = cur / "dataset" @@ -238,10 +241,25 @@ def run_once(config_override=None, batch_size=50, max_batches=None, gc.collect() model, dev = infer.load_model(ckpt_path=None) + cuda = (dev.type == "cuda") + + # 本地从已建好的 batches 构造 rep 缓存（复用 batches、省内存；不计入计时） + if want_precompute: + lc, ec = [], [] + with torch.inference_mode(): + for b in batches: + bb = infer.move_batch_to_device(b, dev) + rep = model.rep_encoder(bb) + lc.append(bb["logid"].to(dev)) + ec.append(rep) + logids = torch.cat(lc) + emb = torch.cat(ec) + order = torch.argsort(logids) + model._rep_cache = (logids[order].contiguous(), emb[order].contiguous()) + print(f"[BENCH] rep cache built from batches: {logids.numel()} items") logid2p = {} t_sum = 0.0 - cuda = (dev.type == "cuda") with torch.inference_mode(): for b in batches: b = infer.move_batch_to_device(b, dev) @@ -300,6 +318,8 @@ def _parse_args(): ap.add_argument("--emb-fp16", action="store_true", help="Embedding表转FP16(查表带宽减半,测AUC)") ap.add_argument("--dedup-emb", action="store_true", help="查表前对sign去重(减少大表随机访存)") ap.add_argument("--sparse-pool", action="store_true", help="稀疏矩阵乘做池化(段内高重复时省)") + ap.add_argument("--precompute-rep", action="store_true", + help="预计算RepEncoder缓存,model(batch)跳过embedding层(从batches自建)") ap.add_argument("--profile", type=int, default=None, metavar="N", help="剖析前 N 个 batch，打印按 CUDA 耗时排序的算子表（定位瓶颈）") ap.add_argument("--rebuild", action="store_true", help="强制重建过滤缓存") @@ -337,6 +357,8 @@ if __name__ == "__main__": cfg["dedup_embedding"] = True if a.sparse_pool: cfg["sparse_pool"] = True + if a.precompute_rep: + cfg["precompute_rep"] = True if a.compile: cfg["compile"] = True if a.profile is not None: diff --git a/代码/code/infer.py b/代码/code/infer.py index cce23e0..a5e8c4c 100644 --- a/代码/code/infer.py +++ b/代码/code/infer.py @@ -55,6 +55,8 @@ CONFIG = { "dedup_embedding": True, # True=查表前对sign去重(只查唯一值再展开),本地7.80->6.49s,AUC逐位等价 "sparse_pool": False, # True=用(段×唯一)稀疏矩阵乘做池化,避免materialize整个[M,512](段内高重复时省) "compile": False, # 是否 torch.compile（实测慢5×，勿开） + "precompute_rep": True, # True=不计时的load_model里预计算所有item的RepEncoder向量, + # model(batch)按logid gather缓存、跳过embedding层(逐位等价) } @@ -624,6 +626,19 @@ class CTRModel(nn.Module): self.seq_encoder = seq_encoder self.d_model = d_model self.linear = nn.Linear(d_model, 1) + self._rep_cache = None # (sorted_logids[N], rep_emb[N, d_model]) 或 None + + def _gather_rep(self, batch): + """有预计算缓存时，按 logid gather 出 RepEncoder 向量（跳过 embedding 层）。 + searchsorted+gather 全在 GPU、无同步。任何缺失 logid → 回退现算整个 batch。""" + sorted_logids, rep_emb = self._rep_cache + logids = batch["logid"].to(sorted_logids.device) + rows = torch.searchsorted(sorted_logids, logids) + rows = rows.clamp(max=sorted_logids.numel() - 1) + hit = sorted_logids[rows] == logids + if bool(hit.all()): # 命中全部 → 直接 gather + return rep_emb[rows].to(self.linear.weight.dtype) + return self.rep_encoder(batch) # 有缺失 → 安全回退 def get_sequence_causal_mask(self, seq_info): lengths = seq_info[1:] - seq_info[:-1] @@ -673,7 +688,10 @@ class CTRModel(nn.Module): return create_block_mask(mask_mod, B=None, H=None, Q_LEN=S, KV_LEN=S, device=device) def forward(self, batch): - seq_input = self.rep_encoder(batch) + if self._rep_cache is not None: + seq_input = self._gather_rep(batch) # 用预计算缓存，跳过 embedding 层 + else: + seq_input = self.rep_encoder(batch) user_offsets = batch["user_offsets"] attn = _resolve_attn(seq_input.device) if attn == "chunked": @@ -697,6 +715,38 @@ class CTRModel(nn.Module): return pred_logits, moe_loss +# ============================================================ +# RepEncoder 预计算缓存 +# ============================================================ + +def build_rep_cache(model, item_dict, user_seq, test_logids_ordered, + max_feasign_per_slot, device, batch_users=200): + """预计算所有 item 的 RepEncoder 向量（context-free），按 logid 排序存入 model._rep_cache。 + + 复用 CTRTestSeqDataset + collate + model.rep_encoder，保证与 model(batch) 内的 + RepEncoder 输出逐位一致。注意：必须用与评测端一致的 max_feasign_per_slot（基线为 {1:2}）， + 否则缓存的 item 向量与 batch 实际特征不符。 + """ + ds = CTRTestSeqDataset( + test_logids_ordered=test_logids_ordered, item_dict=item_dict, + user_seq=user_seq, max_feasign_per_slot=max_feasign_per_slot, max_ctx_len=None) + loader = DataLoader(ds, batch_size=batch_users, shuffle=False, num_workers=0, + collate_fn=make_collate_fn(ds.max_slot_id)) + logid_chunks, emb_chunks = [], [] + model.eval() + with torch.inference_mode(): + for batch in loader: + batch = move_batch_to_device(batch, device) + rep = model.rep_encoder(batch) # [num_tokens, d_model] + logid_chunks.append(batch["logid"].to(device)) + emb_chunks.append(rep) + logids = torch.cat(logid_chunks) + emb = torch.cat(emb_chunks) + order = torch.argsort(logids) + model._rep_cache = (logids[order].contiguous(), emb[order].contiguous()) + return model._rep_cache + + # ============================================================ # 模型加载入口 # ============================================================ @@ -779,6 +829,29 @@ def load_model(ckpt_path, device='cuda:0'): print(f"[INFO] attention={_resolve_attn(dev)}, " f"moe={'dense' if CONFIG.get('vectorize_moe', True) else 'loop'}") + # === 预计算 RepEncoder 缓存（不计时阶段）=== + if CONFIG.get("precompute_rep", False) and model._rep_cache is None: + try: + ds_dir = None + for cand in (Path(ckpt_path).parent / "dataset", Path("dataset"), + Path(__file__).parent / "dataset"): + if cand.exists(): + ds_dir = cand + break + if ds_dir is not None: + history = ds_dir / "history" + test_csv = ds_dir / "test.csv" + files = (sorted(history.glob("*.csv")) if history.exists() else []) + [test_csv] + item_dict, user_seq = load_sample_files(files) + test_logids = list(load_logids_from_file(test_csv)) + build_rep_cache(model, item_dict, user_seq, test_logids, {1: 2}, dev) + print(f"[INFO] rep cache built: {model._rep_cache[0].numel()} items") + else: + print("[INFO] dataset/ not found, skip rep precompute (fallback to in-batch)") + except Exception as e: + print(f"[WARNING] rep precompute failed ({e}), fallback to in-batch RepEncoder") + model._rep_cache = None + if CONFIG.get("compile", False): try: model = torch.compile(model, dynamic=True)