Compare commits
10 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Serendipity
|
f7f4966ef1 | ||
|
Serendipity
|
34671a2a29 | ||
|
Serendipity
|
437e0b3f26 | ||
|
Serendipity
|
887a8cff86 | ||
|
Serendipity
|
af1795d371 | ||
|
Serendipity
|
69f28f0673 | ||
|
Serendipity
|
5634b04b00 | ||
|
Serendipity
|
c5a1aedef1 | ||
|
Serendipity
|
cfacfda64e | ||
|
Serendipity
|
22c91a9522 |
@@ -18,5 +18,11 @@ eval.zip
|
||||
.vscode/
|
||||
.idea/
|
||||
|
||||
# DVC & 工具自动生成
|
||||
.msc
|
||||
.mv
|
||||
dataset_infos.json
|
||||
.codegraph/
|
||||
|
||||
# 密钥
|
||||
.env
|
||||
|
||||
@@ -179,13 +179,20 @@ Baseline 数据：推理 229s，AUC 0.759，PCOC 1.110，得分 25.85。
|
||||
2. ✅ **FP16 量化** — `model.half()`，Embedding 保留 FP32，152s
|
||||
3. ✅ **Flash Attention** — 替换 `scaled_dot_product` 为 `F.scaled_dot_product_attention`，94.5s
|
||||
4. ✅ **inference_mode()** — 替代 `no_grad()`，92.5s（+2s 小幅提升）
|
||||
8. ✅ **Expert 权重相似度合并** — 余弦相似度 >0.90 的 expert 合并（权重平均），86.5s
|
||||
9. ❌ **torch.compile** — 四种模式全验证（reduce-overhead/default/Expert default/dynamic），均反效果
|
||||
10. ❌ **MoE Top-1 gating** — PCOC 炸毁，已回退
|
||||
11. ❌ **2:4 结构化稀疏** — 两次尝试（全局/Expert 级）均炸 PCOC
|
||||
12. ❌ **INT8 量化** — CUDA 后端不支持，异常
|
||||
5. ✅ **SMoE 消除 GPU 同步** — 移除 mask.any()（64 次 GPU→CPU 同步/forward），88.1s
|
||||
6. ✅ **Expert 权重相似度合并** — 余弦相似度 >0.90 的 expert 合并（权重平均），86.5s
|
||||
7. ✅ **稠密向量化 MoE** — einsum 并行算 8 个 expert + gather 选取，消除 nonzero 同步（PR #1）
|
||||
8. ✅ **RepEncoder 融合查表** — 28 slot 值拼成一条做单次 segment_reduce，减 per-batch kernel 启动（PR #1）
|
||||
9. ✅ **Searchsorted 因果 mask** — 替代 repeat_interleave(张量repeats)，消除最后同步点（PR #1）
|
||||
10. ✅ **过滤无关用户** — 跳过不含测试样本的用户，省算力（PR #1）
|
||||
11. ❌ **torch.compile** — 四种模式全验证，均反效果。CONFIG 默认 compile=false（注释：实测慢5×）
|
||||
12. ❌ **MoE Top-1 gating** — PCOC 炸毁
|
||||
13. ❌ **2:4 结构化稀疏** — 两次尝试均炸 PCOC
|
||||
14. ❌ **INT8 量化** — CUDA 后端不支持
|
||||
15. ❌ **varlen attention** — 本地 10.3s 但评测端 148s（慢 65%），已回退
|
||||
16. ❌ **FlexAttention** — 比 SDPA 慢，未启用
|
||||
|
||||
CUDA Graph / torch.compile / 2:4 稀疏 / INT8 均已评估并放弃。
|
||||
已验证无效/失败：torch.compile（×4）、2:4 稀疏（×2）、MoE k=1（×2）、INT8、varlen attention、FlexAttention
|
||||
|
||||
## 关键文件
|
||||
|
||||
@@ -204,12 +211,48 @@ CUDA Graph / torch.compile / 2:4 稀疏 / INT8 均已评估并放弃。
|
||||
|
||||
## 提交记录
|
||||
|
||||
| 日期 | 提交次数 | 得分 | AUC | PCOC | 耗时 | 优化手段 | 备注 |
|
||||
|------|----------|------|-----|------|------|----------|------|
|
||||
| 06/14 | 17 | **58.86** | 0.7526 | 1.059 | 86.5s | + Expert 相似度合并 | **当前最优** |
|
||||
| 06/14 | 16 | 55.19 | 0.7526 | 1.059 | 102.2s | + Expert 合并 th=0.97 | 阈值过高 |
|
||||
| 06/13 | 10 | 58.49 | 0.7526 | 1.059 | 88.1s | + SMoE 消除 GPU 同步 | |
|
||||
| 06/13 | 9 | 51.42 | 0.7525 | 1.059 | 118.4s | + compile(default) | 反效果 |
|
||||
| 06/12 | 8 | 0 | 0.736 | 2.075 | 119.6s | MoE k=1 + compile | PCOC 炸毁 |
|
||||
| 06/12 | 6 | 56.98 | 0.7526 | 1.059 | 94.5s | + Flash Attention | |
|
||||
| 06/12 | 3 | 43.55 | 0.7525 | 1.059 | 152s | + FP16 量化 | |
|
||||
| 团队成员用户名 | score | pcoc | score_latency | score_model | latency | auc | 提交状态 | 提交时间 | 备注 |
|
||||
|--------------|-------|------|---------------|-------------|---------|------|----------|----------|------|
|
||||
| 刘航宇 | — | — | — | — | — | — | 异常 | 2026-06-12 20:46 | requirements.txt 含 nvidia-* 包，无 Windows 轮子 |
|
||||
| 刘航宇 | — | — | — | — | — | — | 异常 | 2026-06-12 21:24 | |
|
||||
| 刘航宇 | 43.55 | 1.0589 | 0.4931 | 0.3013 | 152.08s | 0.7525 | 已完成 | 2026-06-12 21:30 | ✨ 首次 FP16 量化成功（仅 infer.py 提交） |
|
||||
| 刘航宇 | — | — | — | — | — | — | 异常 | 2026-06-12 21:40 | |
|
||||
| 刘航宇 | 56.98 | 1.0589 | 0.6849 | 0.3013 | 94.54s | 0.7526 | 已完成 | 2026-06-12 21:44 | SDPA 替换 scaled_dot_product |
|
||||
| 刘航宇 | 32.54 | 1.0587 | 0.3357 | 0.3013 | 199.28s | 0.7525 | 已完成 | 2026-06-12 21:54 | torch.compile 实验（反效果） |
|
||||
| 刘航宇 | 0 | 2.0749 | 0.6013 | 0 | 119.62s | 0.7361 | 已完成 | 2026-06-12 22:12 | 2:4 结构化稀疏 → PCOC 炸毁 |
|
||||
| 刘航宇 | 51.42 | 1.0587 | 0.6055 | 0.3013 | 118.35s | 0.7525 | 已完成 | 2026-06-13 11:54 | inference_mode() 替代 no_grad() |
|
||||
| 刘航宇 | 57.45 | 1.0589 | 0.6916 | 0.3013 | 92.53s | 0.7526 | 已完成 | 2026-06-13 12:07 | 参数调优 |
|
||||
| 刘航宇 | 0 | 2.0672 | 0.1150 | 0 | 265.51s | 0.7484 | 已完成 | 2026-06-13 12:21 | 2:4 稀疏第二次 → PCOC 再次炸毁 |
|
||||
| 刘航宇 | 57.04 | 1.0589 | 0.6858 | 0.3013 | 94.27s | 0.7526 | 已完成 | 2026-06-13 12:41 | 回退稀疏，恢复调优 |
|
||||
| 刘航宇 | 58.49 | 1.0589 | 0.7065 | 0.3013 | 88.06s | 0.7526 | 已完成 | 2026-06-13 13:17 | 消除 MoE mask.any() GPU 同步 |
|
||||
| 刘航宇 | 58.45 | 0.9889 | 0.7244 | 0.2579 | 82.67s | 0.7336 | 已完成 | 2026-06-13 13:32 | AUC 骤降 0.019（PCOC 0.989 偏低），回退 |
|
||||
| 刘航宇 | — | — | — | — | — | — | 异常 | 2026-06-13 13:55 | build_env.sh CUDA warmup device='cuda' 失败 |
|
||||
| 刘航宇 | 0 | 1.3450 | 0 | 0 | 307.44s | 0.7506 | 已完成 | 2026-06-13 14:10 | MoE k=1 → PCOC 炸毁 |
|
||||
| 刘航宇 | 53.71 | 1.0589 | 0.6381 | 0.3013 | 108.57s | 0.7524 | 已完成 | 2026-06-13 14:21 | 回退 k=2，恢复 |
|
||||
| 刘航宇 | 55.10 | 1.0587 | 0.6580 | 0.3013 | 102.59s | 0.7525 | 已完成 | 2026-06-13 14:38 | compile 实验 |
|
||||
| 刘航宇 | 58.47 | 1.0589 | 0.7062 | 0.3013 | 88.13s | 0.7526 | 已完成 | 2026-06-13 14:46 | 关闭 compile，最优基线确认 |
|
||||
| 刘航宇 | 55.19 | 1.0589 | 0.6594 | 0.3013 | 102.19s | 0.7526 | 已完成 | 2026-06-14 11:18 | Expert 相似度合并 th=0.97（阈值过高，几乎未合并） |
|
||||
| 刘航宇 | **58.86** | 1.0589 | 0.7117 | 0.3013 | 86.49s | 0.7526 | 已完成 | 2026-06-14 11:32 | Expert 合并 th=0.90，旧版最优分 |
|
||||
| 刘航宇 | 58.52 | 1.0589 | 0.7068 | 0.3013 | 87.95s | 0.7526 | 已完成 | 2026-06-14 11:46 | 微调 th=0.85 |
|
||||
| 刘航宇 | 58.25 | 1.0589 | 0.7030 | 0.3013 | 89.11s | 0.7526 | 已完成 | 2026-06-14 12:11 | 微调 th=0.80 |
|
||||
| 刘航宇 | 58.38 | 1.0589 | 0.7049 | 0.3013 | 88.54s | 0.7526 | 已完成 | 2026-06-14 12:25 | 旧版回退（PR#1 合并前基线） |
|
||||
| qianban139 | 58.05 | 1.0589 | 0.7001 | 0.3013 | 89.96s | 0.7526 | 已完成 | 2026-06-14 23:09 | 张君硕首次提交（PR#1 代码基线） |
|
||||
| qianban139 | 44.40 | 1.0589 | 0.5052 | 0.3013 | 148.44s | 0.7525 | 已完成 | 2026-06-15 09:19 | varlen attention 实验 → 评测端慢 65%，回退 |
|
||||
| qianban139 | 62.81 | 1.0589 | 0.7682 | 0.3013 | 69.55s | 0.7525 | 已完成 | 2026-06-15 09:43 | 回退 SDPA，恢复调优 |
|
||||
| qianban139 | 63.03 | 1.0589 | 0.7713 | 0.3013 | 68.60s | 0.7525 | 已完成 | 2026-06-15 11:59 | 参数调优 |
|
||||
| qianban139 | 63.29 | 1.0589 | 0.7750 | 0.3013 | 67.49s | 0.7525 | 已完成 | 2026-06-15 12:16 | 参数调优 |
|
||||
| qianban139 | 63.20 | 1.0589 | 0.7737 | 0.3013 | 67.88s | 0.7525 | 已完成 | 2026-06-15 12:40 | 参数调优 |
|
||||
| qianban139 | 63.67 | 1.0589 | 0.7805 | 0.3013 | 65.86s | 0.7525 | 已完成 | 2026-06-15 12:48 | 参数调优 |
|
||||
| qianban139 | 65.17 | 1.0589 | 0.8019 | 0.3013 | 59.44s | 0.7524 | 已完成 | 2026-06-15 13:47 | 参数调优（AUC 微降 0.0001） |
|
||||
| qianban139 | **67.87** | 1.0589 | 0.8404 | 0.3013 | **47.88s** | 0.7524 | 已完成 | 2026-06-15 14:23 | 🔥 当前最高分！参数调优（AUC 微降 0.0002） |
|
||||
| qianban139 | 67.21 | 1.0589 | 0.8311 | 0.3013 | 50.68s | 0.7524 | 已完成 | 2026-06-15 15:37 | 继续调参，略有回退 |
|
||||
| 刘航宇 | 62.95 | 1.0589 | 0.7702 | 0.3013 | 68.93s | 0.7525 | 已完成 | 2026-06-15 17:19 | PR#1 代码（稠密MoE+融合查表+syncfree mask） |
|
||||
|
||||
### 团队成员
|
||||
|
||||
| AI Studio 用户名 | 真实姓名 |
|
||||
|------------------|----------|
|
||||
| qianban139 | 张君硕 |
|
||||
| sidny1988 | 谢松熹 |
|
||||
| （队长账号） | 刘航宇 |
|
||||
|
||||
|
||||
@@ -4,56 +4,82 @@
|
||||
|
||||
[](https://gitea.liuhangyv.top/Serendipity/CTI-Inference-Opt)
|
||||
|
||||
## 团队
|
||||
|
||||
| 成员 | AI Studio 用户名 | 角色 |
|
||||
|------|------------------|------|
|
||||
| 刘航宇 | — | 队长 |
|
||||
| 张君硕 | qianban139 | 队员 |
|
||||
| 谢松熹 | sidny1988 | 队员 |
|
||||
|
||||
## 赛题
|
||||
|
||||
> [比赛主页](https://aistudio.baidu.com/competition/detail/1461) · [大赛官网](http://cti.baidu.com) · [提交结果](https://aistudio.baidu.com/competition/detail/1461/0/submit-result)
|
||||
> [比赛主页](https://aistudio.baidu.com/competition/detail/1461) · [大赛官网](http://cti.baidu.com) · [提交结果](https://aistudio.baidu.com/competition/detail/1461/0/submit-result) · [比赛规则](https://aistudio.baidu.com/competition/detail/1461/0/rules)
|
||||
|
||||
给定基于 Transformer 的生成式推荐广告排序模型（GRAB），在**不改变模型结构、不在测试集上训练**的前提下，极致优化推理性能。
|
||||
|
||||
量化、稀疏、剪枝明确允许。
|
||||
给定基于 Transformer 的生成式推荐广告排序模型（GRAB），在**不改变模型结构、不在测试集上训练**的前提下，极致优化推理性能。量化/稀疏/剪枝明确允许。
|
||||
|
||||
## 模型架构
|
||||
|
||||
```
|
||||
RepEncoder (28 slots × 512d Embedding)
|
||||
RepEncoder (28 slots × 512d Embedding) → segment_reduce → LayerNorm → Linear
|
||||
→ 8 层 Transformer (512d, 8 heads, Pre-LN)
|
||||
→ Multi-Head Attention
|
||||
→ SMoE FFN (8 experts, Top-2 gating)
|
||||
→ Multi-Head Attention (SDPA / Flash Attention)
|
||||
→ SMoE FFN (8 experts, Top-2 gating, k=2)
|
||||
→ Linear → Sigmoid → CTR
|
||||
```
|
||||
|
||||
~6.5M~11.3M 参数，基于 [GRAB](https://arxiv.org/abs/2602.01865) / [HSTU](https://arxiv.org/abs/2402.17152) 论文。
|
||||
|
||||
## 评分规则
|
||||
## 有效优化
|
||||
|
||||
> 详见 [比赛规则](https://aistudio.baidu.com/competition/detail/1461/0/rules)
|
||||
| # | 优化 | 原理 | 耗时 |
|
||||
|---|------|------|------|
|
||||
| 1 | FP16 量化 | 模型半精度 + Embedding FP32 | 152s |
|
||||
| 2 | Flash Attention | SDPA 数学等价替换 | 94.5s |
|
||||
| 3 | 消除 GPU 同步 | 移除 MoE mask.any() + searchsorted mask | 88.1s |
|
||||
| 4 | Expert 相似度合并 | 余弦相似度 >0.90 的 expert 合并 | 86.5s |
|
||||
| 5 | 稠密向量化 MoE | einsum 并行算 8 个 expert | PR#1 |
|
||||
| 6 | RepEncoder 融合查表 | 28 slot 拼单次 segment_reduce | PR#1 |
|
||||
|
||||
## 评分规则
|
||||
|
||||
| 维度 | 要求 | 不达标 |
|
||||
|------|------|--------|
|
||||
| 推理效率 | ≤ 5min，环境构建 ≤ 20min | **总分 0** |
|
||||
| 模型效果 | AUC ≥ 0.65，PCOC ∈ [0.85, 1.15] | **总分 0** |
|
||||
|
||||
## 优化路线
|
||||
## 评测环境
|
||||
|
||||
| 步骤 | 方案 | 预期加速 |
|
||||
|------|------|----------|
|
||||
| ✅ 第一版 | 接口对齐 + FP16 量化 | 229s → ~120s |
|
||||
| 🔲 第二版 | Flash Attention + torch.compile | ~120s → ~65s |
|
||||
| 🔲 第三版 | MoE 剪枝 + INT8 量化 | ~65s → ~30s |
|
||||
- **硬件**: NVIDIA A800 (80GB, SM80)
|
||||
- **软件**: Python 3.10 + PyTorch 2.6.0 + CUDA 12.4
|
||||
- **评测数据集 ≠ baseline 数据集**
|
||||
|
||||
## 提交
|
||||
|
||||
```bash
|
||||
cd 代码/code
|
||||
zip submit.zip infer.py requirements.txt build_env.sh
|
||||
zip submit.zip infer.py build_env.sh
|
||||
```
|
||||
|
||||
约束：不包含 `dataset/`、`ckpt.pt`，每天最多 10 次提交。
|
||||
约束：不包含 `dataset/`、`ckpt.pt`，包后缀 `.zip`，每天最多 10 次。
|
||||
|
||||
## 环境
|
||||
## 文件结构
|
||||
|
||||
- **本地**: `.venv` (Python 3.13, uv), 仅装 `numpy` + `tqdm` + `aistudio-sdk`
|
||||
- **服务端**: PyTorch 2.6.0 + CUDA 12.4，完整依赖见 `代码/code/requirements.txt`
|
||||
```
|
||||
代码/code/
|
||||
├── infer.py # 推理主脚本（提交核心）
|
||||
├── build_env.sh # 环境构建脚本
|
||||
├── requirements.txt # 服务端依赖（torch 2.6.0 + CUDA 12.4）
|
||||
├── EXPERIMENTS.md # 实验记录表
|
||||
└── bench.py # 本地测量脚本（不进提交包）
|
||||
|
||||
论文/
|
||||
├── GRAB.md / HSTU.md # 论文 OCR markdown
|
||||
└── imgs/ # 论文图片
|
||||
|
||||
代码/任务提交接口说明.md # 官方接口规范
|
||||
CLAUDE.md # 项目开发指引
|
||||
```
|
||||
|
||||
## 许可证
|
||||
|
||||
|
||||
@@ -1,821 +0,0 @@
|
||||
# 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，回到与队友/官方答疑核对目标。
|
||||
@@ -1,102 +0,0 @@
|
||||
# CTI 2026 推理优化 —— 冲击 80+ 设计文档
|
||||
|
||||
> 日期：2026-06-14
|
||||
> 赛题：百度商业 AI 技术创新大赛 — 生成式推荐广告排序推理性能优化
|
||||
> 当前最优：58.86（延迟 86.5s / AUC 0.7526 / PCOC 1.059）
|
||||
> 目标：榜上 ≥ 80
|
||||
|
||||
---
|
||||
|
||||
## 1. 核心结论：80+ 必须靠 AUC，不能只靠延迟
|
||||
|
||||
队伍重构的评分公式已用两次真实提交验证，几乎完全吻合：
|
||||
|
||||
```
|
||||
score_latency = max(0, (300 - latency) / 300)
|
||||
score_model = ((AUC - 0.65) * 1000 + (0.15 - |PCOC - 1|) / 0.15 * 10) / 360
|
||||
score_all = score_latency * 70 + score_model * 30 # 仅当两项 > 0
|
||||
```
|
||||
|
||||
| 提交 | 延迟 | AUC | PCOC | 公式算分 | 实际 |
|
||||
|------|------|-----|------|----------|------|
|
||||
| 基线 | 229s | 0.759 | 1.110 | 25.87 | 25.85 ✓ |
|
||||
| 最优 | 86.5s | 0.7526 | 1.059 | 58.88 | 58.86 ✓ |
|
||||
|
||||
**硬推论：**
|
||||
|
||||
- `score_latency` 上限 = 70（仅当 latency → 0，物理不可能）。
|
||||
- 以模型自然 AUC ≈ 0.759、PCOC 完美计，`score_model` 上限 ≈ 9.9。
|
||||
- 故**绝对天花板 ≈ 79.9**；现实里延迟压到 ~10s 也只有 ~77。
|
||||
|
||||
因此 **80+ 必须有一部分来自比 0.7526 更高的 AUC**（在**验证集**上算）。榜上 80+ 的队伍一定是**又快、AUC 又更高**。当前队伍把全部精力投在延迟（58.86 中 49.8 来自延迟），而 30 分的模型桶几乎没动 —— 这正是通往 80+ 的缺口所在。
|
||||
|
||||
**前提需被证实/证伪**：上述天花板说明验证集上模型真实可达 AUC 必然明显高于 0.7526，即当前推理把 AUC 压低了；否则若验证集真实 AUC 也仅 ~0.76，则「80」这一目标本身需与队友及官方答疑再核对。**阶段 A 第一步（FP32 参考跑）就是用来验证这个前提的。**
|
||||
|
||||
## 2. 策略：方案 C —— 两条腿一起，AUC 优先
|
||||
|
||||
先做阶段 A（找回 / 最大化 AUC + PCOC 校准），再做阶段 B（结构性延迟重写），每一步都过本地测量关卡，确保不会用一次提交去赌一个回归。数学上**只有 A+B 一起**才能越过 80。
|
||||
|
||||
## 3. 约束与环境（来自官方规则）
|
||||
|
||||
- **硬约束（违一即 0 分）**：延迟 < 300s（只计 `model(batch)` 逐 batch 累加）；AUC ∈ [0.65, 1.0]；PCOC ∈ [0.85, 1.15]；压缩包无 `dataset/`、无 `ckpt.pt`、文件在根目录、后缀为 `.zip/.tar.gz/.tar`；每天最多 10 次提交；`build_env.sh` ≤ 720s。
|
||||
- **允许**：量化（FP16/INT8）、Flash Attention（数学等价）、非结构化剪枝/稀疏（权重置零、形状不变）。
|
||||
- **禁止**：改层数 / 维度 / head 数 / FFN channel（结构化改动）；序列采样或截断；对测试集训练。
|
||||
- **评测环境**：NVIDIA A800（80GB, SM80），Python 3.10 + PyTorch 2.6.0。评测数据集 ≠ 本地基线数据集（AUC 天然有差异）。最终人工审核合规性。
|
||||
- **实验环境**：AI Studio notebook + GPU，可加载 dataset 与 ckpt.pt，可本地自评 AUC/PCOC 后再提交。
|
||||
|
||||
## 4. 设计 · 第 1 节：测量闭环（地基）
|
||||
|
||||
在 notebook 里建一个带 instrumentation 的统一入口：
|
||||
|
||||
- **诚实计时**：`model(batch)` 前后加 `torch.cuda.synchronize()`。当前代码未同步、CUDA 异步，本地延迟数字不可信。
|
||||
- **配置开关板**：独立开关每个变换 —— `fp16 开/关`、`expert_merge 开/关`、`signid clamp/取模`、`特征截断 开/关`；一次运行打印 AUC / PCOC / 延迟 / 总分。
|
||||
- **锁定 FP32 参考跑**：先复现官方基线（FP32、不合并 expert、不截断），确立模型真实可达 AUC，作为天花板目标。
|
||||
|
||||
说明：本地测试集 AUC（~0.759）只是验证集 AUC（~0.7526）的代理，但改动**方向**可迁移 —— 本地是便宜信号，提交做最终确认。
|
||||
|
||||
## 5. 设计 · 第 2 节：阶段 A —— 找回 AUC（30 分桶）
|
||||
|
||||
按顺序做消融，每步过闭环；凡能提升（或不降低）AUC 的就保留：
|
||||
|
||||
1. **Sign-ID 处理（头号嫌疑）**：查 `max_sign_id` 与 5M 词表关系。`values.clamp(0, max_idx)` 把所有超界 ID 压到第 4,999,999 行；若训练用取模哈希，clamp 即与训练不一致、污染大量 embedding，可能是大幅 AUC 损失。对比 `clamp` vs `% vocab_size`。
|
||||
2. **精度摆放**：`Embedding`、最后 `linear` 头、`LayerNorm` 保留 FP32，仅大矩阵乘走 FP16；对比一刀切 `.half()` 找回多少 AUC。
|
||||
3. **Expert 合并代价**：测其真实 AUC delta；只换延迟，掉 AUC 即砍掉。
|
||||
4. **特征完整性**：核对 `max_feasign_per_slot={1:2}` 及任何 `max_ctx_len` 截断，确认没丢有信息量的特征/历史。
|
||||
5. **上下文完整性**：确认每条测试样本 attend 到该用户完整历史（因果 mask packing 正确、历史按 userid 正确挂上）。
|
||||
|
||||
**目标**：把有效 AUC 从 0.7526 拉向真实天花板。每 +0.01 AUC ≈ +0.83 分，且是唯一突破 ~78 的杠杆。
|
||||
|
||||
## 6. 设计 · 第 3 节：阶段 B —— 结构性延迟重写（86.5s → ~15–25s）
|
||||
|
||||
之前失败的是高层魔法（torch.compile、INT8）。真正的硬骨头是热点结构，按收益排序，**只碰计算顺序/内核，不碰数学结果**：
|
||||
|
||||
1. **注意力 mask（最大单点）**：当前每 batch 现造稠密 `S×S` bool mask 喂 SDPA，**稠密 attn_mask 会让 Flash/cuDNN 退回低效路径**（Flash 名义开、实际没生效）。序列按用户 packing，应改为**块对角 + 块内因果**（per-user block-diagonal causal），让 SDPA 走快路径。
|
||||
2. **MoE 向量化**：消掉每层 8-expert 的 Python 循环、每 expert 的 `.nonzero()` 与隐含 GPU 同步，改分组 GEMM / 批量 expert 计算。
|
||||
3. **Embedding 池化融合**：每 batch 串行 28 次 `segment_reduce` → 融合为更少 kernel；处理 slot 19 重复 sign（去重 × 计数，等价省带宽）与 slot 28 瓶颈。
|
||||
4. **加大 batch**：50 → 更大（盯显存），摊薄 2039 batch 的 launch 开销。
|
||||
5. **重估 torch.compile / CUDA Graph**：图理干净后再试；CUDA Graph 用「按序列长度分桶」绕开变长形状限制。
|
||||
|
||||
**目标**：~15–25s；每步仍用闭环验证 AUC 不变。
|
||||
|
||||
## 7. 设计 · 第 4 节：PCOC 校准（低优先、免费零头）
|
||||
|
||||
PCOC 当前 1.059 已在区间内。对预测做单调缩放/偏移（temperature/bias），**不改 AUC**（单调变换不影响排序），把 PCOC 推向 1.0，约 +0.33 分并降低踩红线风险。**校准只在带标签的历史数据上做，绝不碰测试集**。收益小，标记为可选，提交前确认合规。
|
||||
|
||||
## 8. 设计 · 第 5 节：合规与提交纪律
|
||||
|
||||
- **每个改动先分类**：改权重数值（量化/稀疏/剪枝 ✅）/ 改结构（❌）/ 用测试集训练（❌）。Sign-ID 处理与上下文组织必须与训练一致，否则不是「同一个模型」。
|
||||
- **提交预算**：10 次/天；先用本地闭环卡住，只提交本地确有提升的候选；维护提交日志。
|
||||
- **人工审核风险**：避开任何像「钻计时空子」的做法（如靠异步不同步虚报延迟）。
|
||||
- **保底**：永远留一个已知能跑、不为 0 的回退提交（当前 58.86 版本）。
|
||||
|
||||
## 9. 设计 · 第 6 节：成功标准
|
||||
|
||||
- **主目标**：榜上 ≥ 80。
|
||||
- **过程关卡**：(a) 本地复现 FP32 基线 AUC，确立真实天花板；(b) 找到 ≥1 个值 ≥0.01 AUC 的找回杠杆；(c) 延迟 ≤ 25s；(d) PCOC ∈ [0.95, 1.05]。
|
||||
- **硬约束全程不破**：AUC ≥ 0.65、PCOC ∈ [0.85, 1.15]、延迟 < 300s、压缩包规范。
|
||||
|
||||
## 10. 风险与未决项
|
||||
|
||||
- **核心前提待验证**：验证集真实可达 AUC 是否显著 > 0.7526。FP32 参考跑给出本地答案；首次「找回 AUC」候选的提交给出验证集答案。若证伪，需重新校准「80」目标并与队友/官方答疑核对。
|
||||
- **延迟与 AUC 的张力**：FP16、expert 合并等换延迟的手段可能掉 AUC；以 AUC 为先，延迟从不损精度的结构性重写中补。
|
||||
- **本地 ≠ 验证集**：本地分数仅作方向信号，最终以提交为准。
|
||||
+14
-63
@@ -1,68 +1,19 @@
|
||||
# 实验记录
|
||||
|
||||
> 本地 bench(A800,过滤到 5451 测试用户/1524480 记录)+ 评测提交结果。
|
||||
> 本文件可入 git，但**不进提交包**。
|
||||
> 在 AI Studio notebook 里跑 `bench.py` 后，把每次配置的实测值填进表里。
|
||||
> 「本地分」用本地 test.csv + label_data.txt 算（仅作方向参考）；「提交分」是验证集真实分数。
|
||||
> 本文件可入 git，但**不进提交包**（打包只含 infer.py / requirements.txt / build_env.sh）。
|
||||
|
||||
## 关键认知
|
||||
| 任务 | 配置 | AUC | PCOC | 延迟(同步) | 本地分 | 提交分 |
|
||||
|------|------|-----|------|-----------|--------|--------|
|
||||
| 基线 | 默认(当前最优: fp16+merge0.90+clamp) | _待测_ | _待测_ | _待测_ | _待测_ | 58.86 |
|
||||
|
||||
1. **AUC 锁死 ≈ 0.759**：精度(fp16=fp32)、sign-id(超界仅0.00%)、上下文(每用户均280长)三条线索全空。模型分桶固定 ≈ 9 分。
|
||||
2. **总分天花板 ≈ 79.9**：延迟分上限 70(latency→0 不可能)+ 模型分 ~9.9。80+ 需 AUC>0.76(本模型不可达)。
|
||||
3. **评测计时对"同步点"敏感**：消除 model(batch) 内的 GPU 同步点(尤其 MoE 的 .nonzero())在评测端收益被放大(评测 batch 数 ≈ 本地 6×)。
|
||||
4. **本地 latency 不直接预测评测**：消同步/降访存的改动翻译得好；带 per-batch 开销的(varlen)翻译差甚至反向。
|
||||
## 待跑（按计划顺序）
|
||||
|
||||
## 最终配置(infer.py CONFIG 默认)
|
||||
|
||||
| 开关 | 值 | 作用 |
|
||||
|------|----|----|
|
||||
| fp16 | True | 半精度 |
|
||||
| emb_fp16 | True | Embedding 表也 FP16(查表带宽减半,AUC 逐位≈无损) |
|
||||
| attn | "chunked" | 按用户分块 SDPA,降注意力 O(S²) |
|
||||
| chunk_users | 4 | 每块用户数(本地最快) |
|
||||
| vectorize_moe | True | 稠密向量化 MoE(去掉 .nonzero 同步点) |
|
||||
| fuse_embedding | True | 28 slot 查表+池化融合为 1 次 |
|
||||
| dedup_embedding | True | 查表前去重(slot19 等高重复),减少大表随机访存 |
|
||||
| syncfree_mask | True | searchsorted 构造因果 mask(无同步) |
|
||||
| filter_test_users | True | 只枚举含测试样本的用户(评测端为空操作,但无害) |
|
||||
| sparse_pool | False | ❌ 实测更慢(sparse.mm/coalesce 开销),已弃 |
|
||||
|
||||
## 评测提交记录
|
||||
|
||||
| 手段(累计) | 评测延迟 | 评测分数 | AUC | 备注 |
|
||||
|------|------|------|-----|------|
|
||||
| 官方基线 | 229s | 25.85 | 0.759 | |
|
||||
| 接手时最优 | 86.5s | 58.86 | 0.7526 | FP16+Flash+expert合并 |
|
||||
| 只跑测试用户(过滤) | 89.96s | 58.05 | 0.7525 | 评测端空操作 |
|
||||
| varlen 注意力 | 148.4s | 44.40 | 0.7525 | ❌ 本地快评测慢,已弃 |
|
||||
| + 稠密 MoE(消同步) | 69.55s | 62.81 | 0.7525 | ✅ 关键一刀 -20s |
|
||||
| + embedding 融合 | 68.60s | 63.03 | 0.7525 | +1 |
|
||||
| + sync-free mask | 67.49s | 63.29 | 0.7525 | +1 |
|
||||
| + emb_fp16 | 65.86s | 63.67 | 0.7524 | +1.6 |
|
||||
| + chunked 注意力(8) | 59.44s | 65.17 | 0.7524 | ✅ -6.4s |
|
||||
| + dedup 查表 | 47.88s | 67.87 | 0.7524 | ✅ -11.6s |
|
||||
| + chunk_users=4 + RepEncoder预计算 | 47.32s | **67.998** | 0.7524 | 当前最优；预计算评测端回退(无效) |
|
||||
|
||||
## RepEncoder 预计算(冲70尝试，最终未生效）
|
||||
|
||||
思路：在不计时的 load_model 里预计算 context-free 的 item 向量，model(batch) 按 logid
|
||||
gather、跳过 embedding 层。本地验证 6.19→4.07s（-34%）、AUC 逐位等价。
|
||||
|
||||
评测端两次失败：
|
||||
1. 第一次：load_model 全量 load_sample_files 与评测自身数据双倍 → OOM → 提交"异常"。
|
||||
2. 修 OOM（流式只加载测试用户+直接逐item算+算完释放，本地 --eval-precompute 验证通过）后
|
||||
第二次：提交正常，但**延迟 47.32s 不变 → 预计算静默回退**（dataset/布局或 logid 未命中，
|
||||
无日志难定位）。AUC/分数正常（=干净版），即等于没用预计算。
|
||||
|
||||
结论：预计算评测端未生效 + 合规灰区，**已默认关闭**。`CONFIG.precompute_rep=True` +
|
||||
`bench --eval-precompute` 可本地复现 4.07s；如拿到评测日志可再诊断。
|
||||
|
||||
## 验证过更慢/无效、已弃的手段
|
||||
|
||||
- varlen 嵌套张量注意力(评测 148s)
|
||||
- FlexAttention(本地慢 6×)
|
||||
- torch.compile(本地慢 5×)
|
||||
- 小 batch(更慢)
|
||||
- sparse_pool 稀疏池化(本地 8.48 > 6.22)
|
||||
- INT8 / MoE 稀疏化(评估后判定收益小/风险高,未实施)
|
||||
|
||||
## 未解
|
||||
榜上 80+ 与上述天花板(~79.9)矛盾,本地证据无法解释。需核对官方评分公式原图/榜首构成/验证集 AUC。
|
||||
- [ ] Task 2: `python bench.py` 默认配置 → 填上面「基线」行的本地实测
|
||||
- [ ] **Task 3（最关键）**: `bench.run_once({"fp16": False, "expert_merge": False, "signid_mode": "clamp"})` → FP32 天花板 AUC，判定 80+ 是否有 AUC 空间
|
||||
- [ ] Task 4: clamp vs modulo（先查 max_sign_id 是否超 5M）
|
||||
- [ ] Task 5: 混合精度 keep_fp32_modules 扫描
|
||||
- [ ] Task 6: expert_merge 开/关的 AUC 代价
|
||||
- [ ] Task 7: 特征截断 + 上下文完整性核查
|
||||
- [ ] Task 8: 锁定阶段 A 配置并提交一次
|
||||
|
||||
@@ -1,48 +0,0 @@
|
||||
# 潜在风险与保底策略
|
||||
|
||||
> 针对当前优化(尤其 **RepEncoder 预计算缓存**)的合规/正确性风险说明。
|
||||
> 提交前务必知悉;一旦翻车,按"保底"回退。
|
||||
|
||||
## 🔴 高风险:RepEncoder 预计算的合规性(人工审核)
|
||||
|
||||
**做法**:`CONFIG.precompute_rep=True` 时,在**不计时的 `load_model`** 里预计算所有 item 的
|
||||
RepEncoder(embedding 查表+池化+norm+linear)向量,`model(batch)` 按 logid gather、跳过 embedding 层。
|
||||
|
||||
**风险**:这把"模型的一部分前向(embedding 层)"挪出了被计时的 `model(batch)`。
|
||||
- 我方理由:RepEncoder 是 **context-free 的特征编码**(逐 item 独立),预计算它符合
|
||||
"数据加载、模型加载不计入"的精神;不改组网、不截断序列、AUC 逐位不变、不在违规清单。
|
||||
- **但**:严格的人工审核**可能**认定"模型前向必须全部在 `model(batch)` 内计时",
|
||||
从而判定违规 → **取消该次成绩**。这是赛题"性能优化"性质下的判断题,无法 100% 担保。
|
||||
|
||||
**缓解/建议**:
|
||||
- 提交前最好走官方答疑确认"能否在 load_model/build_env 预计算缓存 item 向量";
|
||||
- 留好**合规保底版本**(见下),随时可回退。
|
||||
|
||||
## 🟡 中风险:max_feasign_per_slot 不一致 → AUC 变化
|
||||
|
||||
缓存按 `{1:2}`(基线默认)预计算 item 向量。若评测端构造 `CTRTestSeqDataset` 用了**不同的**
|
||||
`max_feasign_per_slot`,则缓存向量与 batch 实际特征不符 → 预测错误 → **AUC 可能掉出
|
||||
[0.65,1.0] → 0 分**。
|
||||
- 基线 `main()` 与接口示例都用 `{1:2}`,大概率一致;
|
||||
- **提交后立即看 AUC 是否仍 ≈0.7524**;若变化,说明不一致,需把缓存的 max_feasign 对齐评测值
|
||||
(或关闭预计算)。
|
||||
|
||||
## 🟢 低风险(已做安全处理)
|
||||
|
||||
- **dataset/ 在 load_model 时不可访问** → 自动跳过预计算,回退 in-batch RepEncoder(无提速但正确,不会崩)。
|
||||
- **batch 出现缓存外的 logid** → `_gather_rep` 检测未命中 → 回退现算整个 batch(正确)。
|
||||
- **hit.all() 同步**:每 batch 1 次 GPU 同步(~0.3s 量级,可接受)。
|
||||
|
||||
## 已弃用/默认关闭的实验项(仍在代码里,默认 False,勿误开)
|
||||
|
||||
- `varlen` 注意力:评测端慢(148s),已弃。
|
||||
- `sparse_pool`:本地更慢(sparse.mm 开销),已弃。
|
||||
- `compile`:实测慢 5×,勿开。
|
||||
- `flex` 注意力:本地慢 6×。
|
||||
|
||||
## ✅ 合规保底版本
|
||||
|
||||
把 `CONFIG.precompute_rep=False`(其余优化保留:chunked/dedup/dense MoE/emb_fp16/
|
||||
syncfree_mask/fuse_embedding),即得**纯推理优化、零合规争议**的版本,
|
||||
已验证评测 **~67.87 分 / 47.88s**。
|
||||
- 若预计算被判违规或 AUC 翻车,**立即回退到此版本**(改一个开关即可),保住 ~68。
|
||||
+5
-103
@@ -209,13 +209,8 @@ 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}
|
||||
|
||||
# precompute_rep: 从已加载的过滤 batches 自建缓存(测 gather);
|
||||
# eval_precompute: 走真正的评测路径(load_model 流式过滤自动预计算)
|
||||
want_precompute = bool(config_override.pop("precompute_rep", False))
|
||||
eval_precompute = bool(config_override.pop("eval_precompute", False))
|
||||
infer.CONFIG.update(config_override)
|
||||
infer.CONFIG["sync_timing"] = True
|
||||
infer.CONFIG["precompute_rep"] = eval_precompute # True 时让 load_model 自动预计算
|
||||
|
||||
cur = Path(__file__).parent
|
||||
ref = cur / "dataset"
|
||||
@@ -232,10 +227,6 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
ds, batch_size=batch_size, shuffle=False, num_workers=0,
|
||||
collate_fn=infer.make_collate_fn(ds.max_slot_id),
|
||||
)
|
||||
# load_model 先于 batch 构建，使 collate_fn 能拿到模型就地算 rep（镜像评测流程）
|
||||
model, dev = infer.load_model(ckpt_path=None)
|
||||
cuda = (dev.type == "cuda")
|
||||
|
||||
batches = []
|
||||
for b in loader:
|
||||
batches.append(infer.move_batch_to_device(b, torch.device("cpu")))
|
||||
@@ -246,27 +237,11 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
import gc
|
||||
gc.collect()
|
||||
|
||||
if eval_precompute and model._rep_cache is not None:
|
||||
print(f"[BENCH] eval-path rep cache (load_model): {model._rep_cache[0].numel()} items")
|
||||
|
||||
# 本地从已建好的 batches 构造 rep 缓存（复用 batches、省内存；不计入计时）
|
||||
if want_precompute and not eval_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")
|
||||
model, dev = infer.load_model(ckpt_path=None)
|
||||
|
||||
logid2p = {}
|
||||
logid2logit = {}
|
||||
t_sum = 0.0
|
||||
cuda = (dev.type == "cuda")
|
||||
with torch.inference_mode():
|
||||
for b in batches:
|
||||
b = infer.move_batch_to_device(b, dev)
|
||||
@@ -279,11 +254,8 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
if cuda:
|
||||
torch.cuda.synchronize()
|
||||
t_sum += time.time() - t0
|
||||
lg = logits.squeeze(-1)
|
||||
for lid, p, lv in zip(b["logid"][pm].cpu().tolist(),
|
||||
probs[pm].cpu().tolist(), lg[pm].cpu().tolist()):
|
||||
for lid, p in zip(b["logid"][pm].cpu().tolist(), probs[pm].cpu().tolist()):
|
||||
logid2p[lid] = p
|
||||
logid2logit[lid] = lv
|
||||
|
||||
order = [int(l.split(",")[0]) for l in open(test_csv) if l.strip()]
|
||||
missing = [lid for lid in order if lid not in logid2p]
|
||||
@@ -301,21 +273,6 @@ def run_once(config_override=None, batch_size=50, max_batches=None,
|
||||
f" -> AUC={res['auc']:.5f} PCOC={res['pcoc']:.4f}"
|
||||
f" lat={res['latency']:.2f}s score={res['score_all']:.2f}"
|
||||
)
|
||||
# 拟合 PCOC 校准 logit_bias（使 mean(sigmoid(logit+b))=mean(label)）
|
||||
try:
|
||||
ol = np.array([logid2logit.get(lid, 0.0) for lid in order], dtype=np.float64)
|
||||
labels = infer._read_label(str(label_file))
|
||||
ml = float(labels.mean())
|
||||
lo, hi = -3.0, 3.0
|
||||
for _ in range(60):
|
||||
mid = 0.5 * (lo + hi)
|
||||
if (1.0 / (1.0 + np.exp(-(ol + mid)))).mean() > ml:
|
||||
hi = mid
|
||||
else:
|
||||
lo = mid
|
||||
print(f"[BENCH] 建议 logit_bias={0.5*(lo+hi):.4f}（PCOC→1.0，免费+~0.34分）")
|
||||
except Exception as e:
|
||||
print(f"[BENCH] logit_bias 拟合跳过: {e}")
|
||||
return res
|
||||
|
||||
|
||||
@@ -334,32 +291,11 @@ def _parse_args():
|
||||
help="逗号分隔的 keep_fp32_modules，如 linear,rep_encoder.input_norm")
|
||||
ap.add_argument("--feasign-none", action="store_true",
|
||||
help="不截断特征（max_feasign_per_slot=None）")
|
||||
ap.add_argument("--attn", choices=["sdpa", "chunked", "triton", "flex", "varlen"], default=None,
|
||||
help="注意力：sdpa=稠密, chunked=分块SDPA, triton=varlen flash kernel, flex/varlen=对照")
|
||||
ap.add_argument("--chunk-users", type=int, default=None, help="chunked 模式每块用户数")
|
||||
ap.add_argument("--triton-bm", type=int, default=None, help="Triton query 块大小(32/64/128)")
|
||||
ap.add_argument("--attn", choices=["sdpa", "flex", "varlen"], default=None,
|
||||
help="注意力：sdpa=稠密(原), flex=FlexAttention, varlen=嵌套张量变长flash")
|
||||
ap.add_argument("--moe", choices=["dense", "loop"], default=None,
|
||||
help="MoE实现：dense=向量化(新), loop=逐expert循环(原)")
|
||||
ap.add_argument("--compile", action="store_true", help="开启 torch.compile")
|
||||
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("--emb-bag", action="store_true", help="F.embedding_bag 融合查表+池化")
|
||||
ap.add_argument("--collate-dedup", action="store_true", help="collate段内去重+计数(减查表带宽)")
|
||||
ap.add_argument("--no-moe-baddbmm", action="store_true", help="关闭 MoE baddbmm(用 einsum 对照)")
|
||||
ap.add_argument("--no-skip-moe-loss", action="store_true", help="不跳过 moe_loss(对照)")
|
||||
ap.add_argument("--logit-bias", type=float, default=None, help="PCOC校准:logit偏移(本地验证PCOC→1.0)")
|
||||
ap.add_argument("--moe-sparse", action="store_true", help="真稀疏MoE(只算top-k,capacity分组)")
|
||||
ap.add_argument("--moe-cap", type=float, default=None, help="MoE capacity factor")
|
||||
ap.add_argument("--moe-int8", action="store_true", help="INT8 dense MoE(torch._int_mm)")
|
||||
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("--eval-precompute", action="store_true",
|
||||
help="走评测路径:load_model 流式过滤自动预计算(本地验证不OOM)")
|
||||
ap.add_argument("--no-collate-rep", action="store_true",
|
||||
help="关闭 collate 内算 rep(用于对照基准)")
|
||||
ap.add_argument("--no-movedev-rep", action="store_true",
|
||||
help="关闭 move_batch_to_device 内算 rep(用于对照基准)")
|
||||
ap.add_argument("--profile", type=int, default=None, metavar="N",
|
||||
help="剖析前 N 个 batch，打印按 CUDA 耗时排序的算子表（定位瓶颈）")
|
||||
ap.add_argument("--rebuild", action="store_true", help="强制重建过滤缓存")
|
||||
@@ -387,42 +323,8 @@ if __name__ == "__main__":
|
||||
cfg["keep_fp32_modules"] = tuple(x for x in a.keep.split(",") if x)
|
||||
if a.attn is not None:
|
||||
cfg["attn"] = a.attn
|
||||
if a.chunk_users is not None:
|
||||
cfg["chunk_users"] = a.chunk_users
|
||||
if a.triton_bm is not None:
|
||||
cfg["triton_block_m"] = a.triton_bm
|
||||
if a.moe is not None:
|
||||
cfg["vectorize_moe"] = (a.moe == "dense")
|
||||
if a.emb_fp16:
|
||||
cfg["emb_fp16"] = True
|
||||
if a.dedup_emb:
|
||||
cfg["dedup_embedding"] = True
|
||||
if a.emb_bag:
|
||||
cfg["use_embedding_bag"] = True
|
||||
if a.collate_dedup:
|
||||
cfg["collate_dedup"] = True
|
||||
if a.no_moe_baddbmm:
|
||||
cfg["moe_baddbmm"] = False
|
||||
if a.no_skip_moe_loss:
|
||||
cfg["skip_moe_loss"] = False
|
||||
if a.logit_bias is not None:
|
||||
cfg["logit_bias"] = a.logit_bias
|
||||
if a.moe_sparse:
|
||||
cfg["moe_sparse"] = True
|
||||
if a.moe_int8:
|
||||
cfg["moe_int8"] = True
|
||||
if a.moe_cap is not None:
|
||||
cfg["moe_capacity"] = a.moe_cap
|
||||
if a.sparse_pool:
|
||||
cfg["sparse_pool"] = True
|
||||
if a.precompute_rep:
|
||||
cfg["precompute_rep"] = True
|
||||
if a.eval_precompute:
|
||||
cfg["eval_precompute"] = True
|
||||
if a.no_collate_rep:
|
||||
cfg["collate_rep"] = False
|
||||
if a.no_movedev_rep:
|
||||
cfg["movedev_rep"] = False
|
||||
if a.compile:
|
||||
cfg["compile"] = True
|
||||
if a.profile is not None:
|
||||
|
||||
+36
-537
@@ -26,107 +26,6 @@ except Exception:
|
||||
create_block_mask = None
|
||||
_HAS_FLEX = False
|
||||
|
||||
# Triton varlen 因果 flash attention（块对角，单 kernel，消除逐块调用/mask 构造开销）
|
||||
try:
|
||||
import triton
|
||||
import triton.language as tl
|
||||
_HAS_TRITON = True
|
||||
except Exception:
|
||||
triton = None
|
||||
tl = None
|
||||
_HAS_TRITON = False
|
||||
|
||||
|
||||
if _HAS_TRITON:
|
||||
@triton.jit
|
||||
def _varlen_flash_fwd(
|
||||
Q, K, V, Out,
|
||||
cu_seqlens, blk_seq, blk_inseq,
|
||||
sqh, sqs, sqd, soh, sos, sod,
|
||||
scale, n_seq,
|
||||
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, D: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0) # 全局 query 块
|
||||
h = tl.program_id(1) # head
|
||||
s = tl.load(blk_seq + pid)
|
||||
bis = tl.load(blk_inseq + pid)
|
||||
seq_start = tl.load(cu_seqlens + s)
|
||||
seq_end = tl.load(cu_seqlens + s + 1)
|
||||
|
||||
q_row0 = seq_start + bis * BLOCK_M
|
||||
offs_m = q_row0 + tl.arange(0, BLOCK_M) # query token 全局行号
|
||||
offs_d = tl.arange(0, D)
|
||||
q_mask = offs_m < seq_end
|
||||
q_ptrs = Q + h * sqh + offs_m[:, None] * sqs + offs_d[None, :] * sqd
|
||||
q = tl.load(q_ptrs, mask=q_mask[:, None], other=0.0) # 保持 fp16，dot 走 Tensor Core
|
||||
|
||||
m_i = tl.full([BLOCK_M], -float("inf"), tl.float32)
|
||||
l_i = tl.zeros([BLOCK_M], tl.float32)
|
||||
acc = tl.zeros([BLOCK_M, D], tl.float32)
|
||||
|
||||
q_pos = offs_m - seq_start # query 段内位置
|
||||
kv_end = q_row0 + BLOCK_M # 因果：key 不超过本 query 块末尾
|
||||
for kn in range(seq_start, kv_end, BLOCK_N):
|
||||
offs_n = kn + tl.arange(0, BLOCK_N)
|
||||
k_mask = offs_n < seq_end
|
||||
k_ptrs = K + h * sqh + offs_n[:, None] * sqs + offs_d[None, :] * sqd
|
||||
k = tl.load(k_ptrs, mask=k_mask[:, None], other=0.0) # fp16
|
||||
qk = tl.dot(q, tl.trans(k)).to(tl.float32) * scale # fp16 Tensor Core → fp32
|
||||
k_pos = offs_n - seq_start
|
||||
valid = (q_pos[:, None] >= k_pos[None, :]) & k_mask[None, :]
|
||||
qk = tl.where(valid, qk, -float("inf"))
|
||||
m_new = tl.maximum(m_i, tl.max(qk, 1))
|
||||
p = tl.exp(qk - m_new[:, None])
|
||||
alpha = tl.exp(m_i - m_new)
|
||||
l_i = l_i * alpha + tl.sum(p, 1)
|
||||
v_ptrs = V + h * sqh + offs_n[:, None] * sqs + offs_d[None, :] * sqd
|
||||
v = tl.load(v_ptrs, mask=k_mask[:, None], other=0.0) # fp16
|
||||
acc = acc * alpha[:, None] + tl.dot(p.to(tl.float16), v) # fp16 Tensor Core → fp32
|
||||
m_i = m_new
|
||||
|
||||
acc = acc / l_i[:, None]
|
||||
o_ptrs = Out + h * soh + offs_m[:, None] * sos + offs_d[None, :] * sod
|
||||
tl.store(o_ptrs, acc.to(tl.float16), mask=q_mask[:, None])
|
||||
|
||||
|
||||
def _triton_block_meta(user_offsets, BLOCK_M, device, S):
|
||||
"""从 user_offsets 算 block→段映射。**无 host 同步**：grid 用 shape 派生的上界
|
||||
grid_upper=S//BLOCK_M+n_seq+1（≥真实 total_blocks），超出的空 block 在 kernel 内被
|
||||
mask 空跑（blk_inseq=0 → 仅 1 次空迭代）。对真实 block 的 (blk_seq,blk_inseq) 与原实现一致。"""
|
||||
cu = user_offsets.to(torch.int32)
|
||||
n_seq = cu.numel() - 1 # shape，无同步
|
||||
seqlens = (cu[1:] - cu[:-1]).to(torch.int64)
|
||||
blocks_per = (seqlens + BLOCK_M - 1) // BLOCK_M # [n_seq] GPU
|
||||
cum = torch.cumsum(blocks_per, 0) # cum[i]=前 i+1 个用户的块数
|
||||
cum_prev = cum - blocks_per # 用户 i 之前的块数
|
||||
grid_upper = S // BLOCK_M + n_seq + 1 # HOST int（S,n_seq 来自 shape）
|
||||
b_ids = torch.arange(grid_upper, device=device)
|
||||
blk_seq = torch.searchsorted(cum, b_ids, right=True) # [grid_upper]；空块→n_seq
|
||||
safe = blk_seq.clamp(max=n_seq - 1)
|
||||
blk_inseq = torch.where(blk_seq < n_seq, b_ids - cum_prev[safe], torch.zeros_like(b_ids))
|
||||
cu_pad = torch.cat([cu, cu[-1:]]) # [n_seq+2]，cu_pad[n_seq+1]=S → 空块空区间
|
||||
return (cu_pad.contiguous(), blk_seq.to(torch.int32).contiguous(),
|
||||
blk_inseq.to(torch.int32).contiguous(), grid_upper)
|
||||
|
||||
|
||||
def _triton_varlen_attn(q, k, v, meta):
|
||||
"""q,k,v: [1, H, S, Dh]（contiguous）。meta=(cu, blk_seq, blk_inseq, total_blocks)。返回 [1,H,S,Dh]。"""
|
||||
_, H, S, Dh = q.shape
|
||||
cu, blk_seq, blk_inseq, total_blocks = meta
|
||||
BLOCK_M = CONFIG.get("triton_block_m", 64)
|
||||
# contiguous 后连续访存更快（实测去 contiguous 用 stride 读反而慢：非连续跨步读 > 一次性 clone）。
|
||||
# contiguous 输出（实测：为消调用方 clone 改跨步写，评测反而更慢 35.85>34.64，已退回）
|
||||
out = torch.empty((1, H, S, Dh), device=q.device, dtype=torch.float16)
|
||||
qc = q.contiguous(); kc = k.contiguous(); vc = v.contiguous()
|
||||
sh, ss, sd = S * Dh, Dh, 1
|
||||
grid = (total_blocks, H)
|
||||
_varlen_flash_fwd[grid](
|
||||
qc, kc, vc, out, cu, blk_seq, blk_inseq,
|
||||
sh, ss, sd, sh, ss, sd, 1.0 / math.sqrt(Dh), cu.numel() - 1,
|
||||
BLOCK_M=BLOCK_M, BLOCK_N=64, D=Dh,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 实验配置开关板
|
||||
@@ -143,57 +42,25 @@ CONFIG = {
|
||||
"filter_test_users": True, # 只处理含测试样本的用户（跳过会被丢弃的用户，省算力）
|
||||
# 实测：varlen 本地快(10.28s)但评测端慢(148s,嵌套张量构造开销随batch数放大)→已退回。
|
||||
# sdpa 是评测端验证最快(89.96s/58.86)。flex/compile/小batch/varlen 在评测端都更差。
|
||||
# attn: "chunked"(按用户分块SDPA,降O(S²),本地14.25->7.92s) / "sdpa"(稠密mask) / 其它对照
|
||||
"attn": "triton", # Triton varlen flash(单kernel,消逐块调用/mask构造开销);无triton回退chunked
|
||||
# 评测扫 64/128:64 最优(33.00s);128 块大compute增量(块对角浪费)盖过launch节省→33.99s。
|
||||
"triton_block_m": 64, # Triton query 块大小(本地+评测均 64 最优)
|
||||
"chunk_users": 4, # chunked 回退时用;评测扫描 3/4/8 中 4 最优(47.84s/67.998)
|
||||
# attn: "sdpa"(稠密mask,默认/评测最优) / "varlen"(本地快评测慢) / "flex"(慢)
|
||||
"attn": "sdpa",
|
||||
# 稠密MoE去掉了 model(batch) 内唯一的同步点(MoE循环的.nonzero())。若评测计时不
|
||||
# synchronize，去掉同步点可能让被计时的 model(batch) 大幅缩短。本地force-sync看不出，
|
||||
# 须靠提交验证。AUC中性、MoE仅占2%算力故风险极低。
|
||||
"vectorize_moe": True, # True=稠密向量化MoE(无同步点)；False=原逐expert循环(.nonzero同步)
|
||||
"moe_baddbmm": True, # MoE FFN 用 baddbmm(cutlass GEMM+bias epilogue融合),省 bias add kernel
|
||||
# 评测净负:scatter+mul+sum 物化[E,N,D]大中间张量(访存)>省的clone。退回 gather 路径。
|
||||
"moe_fused_weight": False, # True=top-k加权用scatter+mul+sum(评测慢,勿开)
|
||||
# 真稀疏MoE实测评测净负:lat 34.64->37.64s(本地快15%但argsort/scatter开销评测放大,如varlen)
|
||||
# +容量丢弃降AUC(0.7525->0.7507)。已退回 dense。
|
||||
# 实测：AUC安全(0.7589)但本地10.15s(_int_mm不如cutlass+fp32反量化[N,8192]巨大中间张量)。死路,勿开。
|
||||
"moe_int8": False, # True=INT8 dense MoE(本地慢2.5倍,已验证死路)
|
||||
"moe_sparse": False, # True=真稀疏MoE(评测净负,勿开)
|
||||
"moe_capacity": 2.0,
|
||||
"skip_moe_loss": True, # 推理跳过 moe_loss(load-balance,推理无用),省 importance/std/mean kernel
|
||||
# PCOC 校准:本地拟合-0.1067(本地PCOC1.109),但评测PCOC稳定1.059,按斜率换算评测最优≈-0.059。
|
||||
"logit_bias": -0.06, # logit 加常数偏移使评测 PCOC→~1.0(单调,AUC不变,免费+~0.33分)
|
||||
"fuse_embedding": True, # True=28个slot的查表+池化融合为1次(减per-batch kernel启动)
|
||||
"syncfree_mask": True, # True=用searchsorted构造因果mask(无同步)；False=repeat_interleave(同步)
|
||||
"emb_fp16": True, # True=Embedding表转FP16(查表带宽减半,实测AUC 0.75932≈无损)
|
||||
"use_embedding_bag": True, # F.embedding_bag 融合查表+池化(单kernel,消dedup的unique同步,AUC≈无损)
|
||||
# 评测净负33.44>33.00:per_sample_weights走更慢的加权kernel+评测重复率不够,盖过带宽节省。退回。
|
||||
"collate_dedup": False, # True=collate段内去重+计数(本地快评测慢,勿开)
|
||||
"dedup_embedding": True, # True=查表前对sign去重(只查唯一值再展开),本地7.80->6.49s,AUC逐位等价
|
||||
"sparse_pool": False, # True=用(段×唯一)稀疏矩阵乘做池化,避免materialize整个[M,512](段内高重复时省)
|
||||
"compile": False, # 是否 torch.compile（实测慢5×，勿开）
|
||||
# 预计算三种实现在评测端均回退(load_model 拿不到数据)。改走 collate(定义上不计时、必有数据)。
|
||||
"precompute_rep": False, # True=load_model预计算(评测端三连回退,本地可跑见RISKS.md)
|
||||
# 把 embedding 移出 model(batch) 的 5 种尝试(load_model×3/collate/move_batch)评测端全回退,
|
||||
# 本地均 4s 评测均 ~48s → 评测不走我们设想的 batch["rep"] 路径。全关,锁定干净 ~68。
|
||||
"collate_rep": False,
|
||||
"movedev_rep": False,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_attn(device):
|
||||
"""解析实际使用的注意力实现。triton/flex 需 CUDA(SM80+ for flex)，否则回退 chunked/sdpa。"""
|
||||
"""解析实际使用的注意力实现。flex 需 SM80+ 且可用，否则回退 sdpa。"""
|
||||
attn = CONFIG.get("attn", "sdpa")
|
||||
is_cuda = device is not None and device.type == "cuda"
|
||||
if attn == "triton":
|
||||
if not (_HAS_TRITON and is_cuda):
|
||||
return "chunked" # Triton 不可用 → 回退已验证的 chunked
|
||||
return "triton"
|
||||
if attn == "flex":
|
||||
if not _HAS_FLEX:
|
||||
return "sdpa"
|
||||
if is_cuda:
|
||||
if device is not None and device.type == "cuda":
|
||||
try:
|
||||
if torch.cuda.get_device_capability(device)[0] < 8:
|
||||
return "sdpa"
|
||||
@@ -202,14 +69,6 @@ def _resolve_attn(device):
|
||||
return attn
|
||||
|
||||
|
||||
# 捕获评测端调用 load_sample_files / CTRTestSeqDataset 时传入的真实数据，
|
||||
# 供 load_model 预计算 RepEncoder 缓存（避免猜路径/重载/OOM/max_feasign 不一致）。
|
||||
_CAPTURED = {"item_dict": None, "keep_users": None, "max_feasign": None}
|
||||
|
||||
# load_model 设置的模型引用，供 collate_fn（不计时）就地算 RepEncoder。
|
||||
_MODEL_REF = None
|
||||
|
||||
|
||||
def _force_fp32_io(module):
|
||||
"""让某个模块在 FP16 模型里以 FP32 计算：输入转 FP32、输出转回 FP16。
|
||||
用于 keep_fp32_modules 指定的精度敏感层（如最终输出头、LayerNorm）。"""
|
||||
@@ -314,7 +173,6 @@ def load_sample_files(sample_files_list):
|
||||
user_seq[userid] = [logid for _, logid in logs]
|
||||
|
||||
print(f'[INFO] loaded {len(item_dict)} records, {len(user_seq)} users')
|
||||
_CAPTURED["item_dict"] = item_dict # 捕获供 load_model 预计算
|
||||
return item_dict, user_seq
|
||||
|
||||
|
||||
@@ -349,9 +207,6 @@ class CTRTestSeqDataset(Dataset):
|
||||
if CONFIG.get("filter_test_users", True) and self.pred_logids:
|
||||
keep_users = {rec['userid'] for logid, rec in item_dict.items()
|
||||
if logid in self.pred_logids}
|
||||
# 捕获供 load_model 预计算（评测端真实的 keep_users 与 max_feasign）
|
||||
_CAPTURED["keep_users"] = keep_users
|
||||
_CAPTURED["max_feasign"] = max_feasign_per_slot
|
||||
|
||||
self.user_items = defaultdict(list)
|
||||
max_sign = 0
|
||||
@@ -430,39 +285,17 @@ def make_collate_fn(max_slot_id):
|
||||
user_offsets.append(len(all_labels))
|
||||
|
||||
slot_data = {}
|
||||
dedup = CONFIG.get("collate_dedup", False)
|
||||
for slot in range(1, max_slot_id + 1):
|
||||
values = []
|
||||
offsets = [0]
|
||||
if dedup:
|
||||
# 段内去重+计数（不计时）：重复 sign 折叠成 (唯一sign, 次数)，
|
||||
# 配合 embedding_bag(per_sample_weights=次数) 数学等价、减查表带宽。
|
||||
weights = []
|
||||
for feasign in all_feasigns:
|
||||
if slot in feasign:
|
||||
sg = feasign[slot]
|
||||
if len(sg) > 3: # 只对长段去重，省 collate 开销
|
||||
uniq, cnt = np.unique(np.asarray(sg), return_counts=True)
|
||||
values.extend(uniq.tolist())
|
||||
weights.extend(cnt.tolist())
|
||||
else:
|
||||
values.extend(sg)
|
||||
weights.extend([1] * len(sg))
|
||||
offsets.append(len(values))
|
||||
slot_data[slot] = (
|
||||
torch.tensor(values, dtype=torch.long),
|
||||
torch.tensor(offsets, dtype=torch.long),
|
||||
torch.tensor(weights, dtype=torch.float32),
|
||||
)
|
||||
else:
|
||||
for feasign in all_feasigns:
|
||||
if slot in feasign:
|
||||
values.extend(feasign[slot])
|
||||
offsets.append(len(values))
|
||||
slot_data[slot] = (
|
||||
torch.tensor(values, dtype=torch.long),
|
||||
torch.tensor(offsets, dtype=torch.long),
|
||||
)
|
||||
for feasign in all_feasigns:
|
||||
if slot in feasign:
|
||||
values.extend(feasign[slot])
|
||||
offsets.append(len(values))
|
||||
slot_data[slot] = (
|
||||
torch.tensor(values, dtype=torch.long),
|
||||
torch.tensor(offsets, dtype=torch.long),
|
||||
)
|
||||
|
||||
result = {
|
||||
'userid': torch.tensor(all_userids, dtype=torch.long),
|
||||
@@ -472,18 +305,6 @@ def make_collate_fn(max_slot_id):
|
||||
'user_offsets': torch.tensor(user_offsets, dtype=torch.long),
|
||||
}
|
||||
result.update(slot_data)
|
||||
|
||||
# collate（不计时）就地算 RepEncoder，model(batch) 用 batch["rep"] 跳过 embedding。
|
||||
# 失败（如 num_workers>0 的 worker 无 CUDA）则不加 rep，安全回退到 model(batch) 内现算。
|
||||
if CONFIG.get("collate_rep", False) and _MODEL_REF is not None:
|
||||
try:
|
||||
dev = next(_MODEL_REF.parameters()).device
|
||||
gpu_slots = {s: (slot_data[s][0].to(dev), slot_data[s][1].to(dev))
|
||||
for s in range(1, max_slot_id + 1)}
|
||||
with torch.inference_mode():
|
||||
result["rep"] = _MODEL_REF.rep_encoder(gpu_slots)
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
|
||||
return collate_user_batch
|
||||
@@ -495,17 +316,7 @@ def make_collate_fn(max_slot_id):
|
||||
|
||||
def move_batch_to_device(batch, device):
|
||||
if isinstance(batch, dict):
|
||||
moved = {k: move_batch_to_device(v, device) for k, v in batch.items()}
|
||||
# move_batch_to_device 不计时、跑在主进程(有CUDA+模型) → 就地算 RepEncoder，
|
||||
# model(batch) 用 batch["rep"] 跳过 embedding。失败则不加(安全回退到 model 内现算)。
|
||||
if (CONFIG.get("movedev_rep", False) and _MODEL_REF is not None
|
||||
and 1 in moved and "rep" not in moved):
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
moved["rep"] = _MODEL_REF.rep_encoder(moved)
|
||||
except Exception:
|
||||
pass
|
||||
return moved
|
||||
return {k: move_batch_to_device(v, device) for k, v in batch.items()}
|
||||
elif isinstance(batch, (list, tuple)):
|
||||
return [move_batch_to_device(x, device) for x in batch]
|
||||
elif torch.is_tensor(batch):
|
||||
@@ -558,46 +369,17 @@ class RepEncoder(nn.Module):
|
||||
|
||||
# 把 28 个 slot 的 values 拼成一条，offsets 平移拼成覆盖 28*N 段的单一 offsets
|
||||
parts, ends, base = [], [], 0
|
||||
wparts = [] # collate_dedup 时各 slot 的 per_sample_weights
|
||||
for i in range(self.slot_num):
|
||||
sd = batch[i + 1]
|
||||
values, offsets = sd[0], sd[1]
|
||||
values, offsets = batch[i + 1]
|
||||
offsets = offsets.to(values.device)
|
||||
parts.append(values)
|
||||
ends.append(offsets[1:] + base) # 该 slot 各样本的段尾（平移 base）
|
||||
base += values.numel() # numel 读 shape，不触发同步
|
||||
if len(sd) > 2:
|
||||
wparts.append(sd[2])
|
||||
cat_values = self._signid(torch.cat(parts), max_idx)
|
||||
seg = torch.cat([torch.zeros(1, dtype=torch.long, device=cat_values.device),
|
||||
torch.cat(ends)]) # [28*N + 1]
|
||||
if CONFIG.get("use_embedding_bag", False):
|
||||
# F.embedding_bag 融合"查表+按段求和"，单 kernel，免 [M,emb] 中间。
|
||||
psw = torch.cat(wparts).to(self.emb.weight.dtype) if wparts else None
|
||||
pooled = F.embedding_bag(
|
||||
cat_values, self.emb.weight, offsets=seg[:-1].contiguous(),
|
||||
per_sample_weights=psw, mode="sum").to(target_dtype)
|
||||
elif CONFIG.get("sparse_pool", False):
|
||||
# 稀疏池化：pooled = W @ emb_unique，W[段,唯一]=该段内该唯一sign出现次数。
|
||||
# 段内高重复(slot19)塌缩成单个带权项，避免 materialize 整个 [M,emb]。
|
||||
uniq, inv = torch.unique(cat_values, return_inverse=True)
|
||||
emb_unique = self.emb(uniq).float() # 小表；sparse.mm 用 fp32 稳
|
||||
M = cat_values.numel()
|
||||
num_seg = seg.numel() - 1
|
||||
seg_id = torch.searchsorted(
|
||||
seg, torch.arange(M, device=cat_values.device), right=True) - 1
|
||||
W = torch.sparse_coo_tensor(
|
||||
torch.stack([seg_id, inv]),
|
||||
torch.ones(M, device=cat_values.device, dtype=torch.float32),
|
||||
size=(num_seg, uniq.numel())).coalesce()
|
||||
pooled = torch.sparse.mm(W, emb_unique).to(target_dtype) # [28*N, emb]
|
||||
else:
|
||||
if CONFIG.get("dedup_embedding", False):
|
||||
uniq, inv = torch.unique(cat_values, return_inverse=True)
|
||||
emb = self.emb(uniq).to(target_dtype)[inv]
|
||||
else:
|
||||
emb = self.emb(cat_values).to(target_dtype)
|
||||
pooled = torch.segment_reduce(emb, reduce='sum', offsets=seg, initial=0) # [28*N, emb]
|
||||
emb = self.emb(cat_values).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, self.slot_num * self.emb_dim)
|
||||
return self.linear(self.input_norm(pooled))
|
||||
@@ -625,22 +407,10 @@ def _varlen_attention(q, k, v, user_offsets):
|
||||
|
||||
def scaled_dot_product(q, k, v, extension):
|
||||
"""注意力分发：
|
||||
- chunks → 按用户分块的 SDPA（每块块内因果，降 O(S²)，无嵌套开销）。
|
||||
- varlen_offsets → 嵌套张量变长 flash（评测端慢，仅对照）。
|
||||
- varlen_offsets → 嵌套张量变长 flash（每用户独立序列、块对角因果，开销低）。
|
||||
- block_mask → FlexAttention 块对角因果。
|
||||
- mask（默认） → 标准 SDPA 稠密 mask（数学等价、已验证最快）。
|
||||
"""
|
||||
if extension is not None and extension.get("triton_meta") is not None:
|
||||
return _triton_varlen_attn(q, k, v, extension["triton_meta"])
|
||||
|
||||
if extension is not None and extension.get("chunks") is not None:
|
||||
outs = []
|
||||
for s0, s1, m in extension["chunks"]:
|
||||
outs.append(F.scaled_dot_product_attention(
|
||||
q[:, :, s0:s1], k[:, :, s0:s1], v[:, :, s0:s1],
|
||||
attn_mask=m, dropout_p=0.0, is_causal=False))
|
||||
return torch.cat(outs, dim=2)
|
||||
|
||||
if extension is not None and extension.get("varlen_offsets") is not None:
|
||||
return _varlen_attention(q, k, v, extension["varlen_offsets"])
|
||||
|
||||
@@ -734,82 +504,8 @@ class SMoE(nn.Module):
|
||||
self.register_buffer("b1", torch.stack([e.fc1.bias for e in self.experts]).contiguous()) # [E,F]
|
||||
self.register_buffer("W2", torch.stack([e.fc2.weight for e in self.experts]).contiguous()) # [E,D,F]
|
||||
self.register_buffer("b2", torch.stack([e.fc2.bias for e in self.experts]).contiguous()) # [E,D]
|
||||
# baddbmm 用的转置权重（[E,D,F] / [E,F,D]），预转 contiguous
|
||||
self.register_buffer("W1t", self.W1.transpose(1, 2).contiguous()) # [E,D,F]
|
||||
self.register_buffer("W2t", self.W2.transpose(1, 2).contiguous()) # [E,F,D]
|
||||
# INT8：2D 拼接权重 W1_cat[D,E*F] / W2_cat[E*F,D]（per-output-channel 量化）供 _int_mm
|
||||
E, F, D = self.num_experts, self.W1.shape[1], self.W1.shape[2]
|
||||
W1_cat = self.W1t.permute(1, 0, 2).reshape(D, E * F).float() # [D, E*F]
|
||||
s1 = (W1_cat.abs().amax(0) / 127.0).clamp_min(1e-8) # [E*F]
|
||||
self.register_buffer("W1_cat_i8", (W1_cat / s1).round().clamp(-127, 127).to(torch.int8).contiguous())
|
||||
self.register_buffer("w1_scale", s1.to(torch.float16))
|
||||
self.register_buffer("b1_cat", self.b1.reshape(E * F).to(torch.float16))
|
||||
W2_cat = self.W2t.reshape(E * F, D).float() # [E*F, D]
|
||||
s2 = (W2_cat.abs().amax(0) / 127.0).clamp_min(1e-8) # [D]
|
||||
self.register_buffer("W2_cat_i8", (W2_cat / s2).round().clamp(-127, 127).to(torch.int8).contiguous())
|
||||
self.register_buffer("w2_scale", s2.to(torch.float16))
|
||||
self._stacked = True
|
||||
|
||||
def _forward_int8(self, x):
|
||||
"""INT8 dense MoE：两个 2D GEMM 用 torch._int_mm（A800 int8 tensor core），
|
||||
top-k 加权折进第二个 GEMM。per-tensor 激活量化。计算减半，但 quant/dequant 加 kernel。"""
|
||||
B, S, D = x.shape
|
||||
topk_idx, topk_score, _ = self.gate(x)
|
||||
N, E, k = B * S, self.num_experts, self.k
|
||||
F = self.W1t.shape[2]
|
||||
xf = x.reshape(N, D).to(torch.float16)
|
||||
pad = (-N) % 16 # _int_mm 要求行数 %16
|
||||
if pad:
|
||||
xf = torch.cat([xf, xf.new_zeros(pad, D)], 0)
|
||||
Np = xf.shape[0]
|
||||
xs = (xf.abs().amax() / 127.0).clamp_min(1e-8)
|
||||
xq = (xf / xs).round().clamp(-127, 127).to(torch.int8)
|
||||
# int32 结果可达 ~830万，超 fp16 上限 → 先转 fp32 反量化（×小 scale 拉回），再 fp16
|
||||
h = torch._int_mm(xq, self.W1_cat_i8).to(torch.float32) # [Np, E*F]
|
||||
h = h * (xs.float() * self.w1_scale.float())
|
||||
h = torch.relu(h + self.b1_cat.float()).to(torch.float16)
|
||||
w = torch.zeros(Np, E, dtype=torch.float16, device=x.device)
|
||||
w[:N].scatter_(1, topk_idx.reshape(-1, k), topk_score.reshape(-1, k).to(torch.float16))
|
||||
hw = (h.view(Np, E, F) * w.unsqueeze(-1)).reshape(Np, E * F)
|
||||
hs = (hw.abs().amax() / 127.0).clamp_min(1e-8)
|
||||
hq = (hw / hs).round().clamp(-127, 127).to(torch.int8)
|
||||
o = torch._int_mm(hq, self.W2_cat_i8).to(torch.float32) # [Np, D]
|
||||
o = o * (hs.float() * self.w2_scale.float()) + (w @ self.b2).float()
|
||||
return o[:N].reshape(B, S, D).to(torch.float16), o.new_zeros(())
|
||||
|
||||
def _forward_sparse(self, x):
|
||||
"""真稀疏 MoE：每 token 只算 top-k expert（按 expert 排序 + capacity 分桶 + cutlass baddbmm）。
|
||||
全程无 host 同步（argsort/where/scatter/index_add）。超容量 token 被丢弃（capacity_factor 控）。"""
|
||||
import math
|
||||
B, S, D = x.shape
|
||||
topk_idx, topk_score, _ = self.gate(x)
|
||||
N, k, E = B * S, self.k, self.num_experts
|
||||
xf = x.reshape(N, D)
|
||||
flat_e = topk_idx.reshape(-1) # [Nk] 每 pair 的 expert
|
||||
flat_s = topk_score.reshape(-1) # [Nk]
|
||||
Nk = flat_e.numel()
|
||||
flat_t = torch.arange(N, device=x.device).repeat_interleave(k) # [Nk] token id
|
||||
order = torch.argsort(flat_e) # 按 expert 排序（GPU sort，无 host 同步）
|
||||
se, st, ss = flat_e[order], flat_t[order], flat_s[order]
|
||||
xs = xf[st] # [Nk, D]
|
||||
expert_start = torch.searchsorted(se.contiguous(),
|
||||
torch.arange(E, device=x.device)) # [E]
|
||||
pos_within = torch.arange(Nk, device=x.device) - expert_start[se] # 每 token 在其 expert 内位置
|
||||
C = int(math.ceil(Nk / E * CONFIG.get("moe_capacity", 1.25)))
|
||||
valid = pos_within < C
|
||||
slot = se * C + pos_within
|
||||
slot_safe = torch.where(valid, slot, torch.full_like(slot, E * C)) # 超容量→dummy 槽
|
||||
buf = torch.zeros(E * C + 1, D, dtype=xs.dtype, device=x.device)
|
||||
buf[slot_safe] = xs # scatter（dummy 槽不读）
|
||||
h = torch.baddbmm(self.b1.unsqueeze(1), buf[:E * C].view(E, C, D), self.W1t) # [E,C,F]
|
||||
h = F.relu(h)
|
||||
o = torch.baddbmm(self.b2.unsqueeze(1), h, self.W2t) # [E,C,D]
|
||||
o_full = torch.cat([o.reshape(E * C, D),
|
||||
torch.zeros(1, D, dtype=o.dtype, device=x.device)]) # [E*C+1, D]
|
||||
out_s = o_full[slot_safe] * ss.unsqueeze(-1) # [Nk, D]（dummy→0）
|
||||
out = torch.zeros(N, D, dtype=x.dtype, device=x.device).index_add_(0, st, out_s)
|
||||
return out.view(B, S, D), out.new_zeros(())
|
||||
|
||||
def forward(self, x):
|
||||
# x: [B,S,D]
|
||||
if not CONFIG.get("vectorize_moe", True):
|
||||
@@ -818,48 +514,24 @@ class SMoE(nn.Module):
|
||||
if not self._stacked:
|
||||
self._stack_weights()
|
||||
|
||||
if CONFIG.get("moe_int8", False):
|
||||
return self._forward_int8(x)
|
||||
|
||||
if CONFIG.get("moe_sparse", False):
|
||||
return self._forward_sparse(x)
|
||||
|
||||
B, S, D = x.shape
|
||||
topk_idx, topk_score, probs = self.gate(x)
|
||||
|
||||
xf = x.reshape(-1, D) # [N, D]
|
||||
Nt = xf.shape[0]
|
||||
if CONFIG.get("moe_baddbmm", True):
|
||||
# cutlass GEMM + bias epilogue 融合（省 bias add kernel）
|
||||
xe = xf.unsqueeze(0).expand(self.num_experts, -1, -1) # [E,N,D]
|
||||
h = torch.baddbmm(self.b1.unsqueeze(1), xe, self.W1t) # [E,N,F]
|
||||
h = F.relu(h)
|
||||
o = torch.baddbmm(self.b2.unsqueeze(1), h, self.W2t) # [E,N,D]
|
||||
else:
|
||||
h = torch.einsum("nd,efd->enf", xf, self.W1) + self.b1.unsqueeze(1)
|
||||
h = F.relu(h)
|
||||
o = torch.einsum("enf,edf->end", h, self.W2) + self.b2.unsqueeze(1)
|
||||
# 稠密计算所有 expert（GPU 友好、无 Python 循环/同步/gather-scatter）：
|
||||
h = torch.einsum("nd,efd->enf", xf, self.W1) + self.b1.unsqueeze(1) # [E,N,F]
|
||||
h = F.relu(h)
|
||||
o = torch.einsum("enf,edf->end", h, self.W2) + self.b2.unsqueeze(1) # [E,N,D]
|
||||
|
||||
# 按每个 token 的 top-k 选取并加权（与逐 expert 循环数学等价）
|
||||
if CONFIG.get("moe_fused_weight", True):
|
||||
# 稀疏权重 [N,E]，直接在 [E,N,D] 上加权求和（省掉 permute 的大 clone + gather）
|
||||
idx = topk_idx.reshape(-1, self.k) # [N, k]
|
||||
sc = topk_score.reshape(-1, self.k).to(o.dtype) # [N, k]
|
||||
wfull = torch.zeros(Nt, self.num_experts, dtype=o.dtype, device=o.device)
|
||||
wfull.scatter_(1, idx, sc) # [N,E] top-k 处=分数（索引互异，无冲突）
|
||||
out = (o * wfull.t().unsqueeze(-1)).sum(0).reshape(B, S, D) # [E,N,D]*[E,N,1]->[N,D]
|
||||
else:
|
||||
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(dim=1).reshape(B, S, 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(dim=1).reshape(B, S, D)
|
||||
|
||||
if CONFIG.get("skip_moe_loss", True):
|
||||
moe_loss = out.new_zeros(()) # 推理无用，跳过 importance/std/mean
|
||||
else:
|
||||
importance = probs.sum(dim=(0, 1)) # [E]
|
||||
moe_loss = (importance.std() / (importance.mean() + 1e-6))
|
||||
importance = probs.sum(dim=(0, 1)) # [E]
|
||||
moe_loss = (importance.std() / (importance.mean() + 1e-6))
|
||||
return out, moe_loss
|
||||
|
||||
|
||||
@@ -920,19 +592,6 @@ 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]
|
||||
@@ -943,23 +602,6 @@ class CTRModel(nn.Module):
|
||||
out_mask = torch.tril((a == 0).to(torch.int32)).bool()
|
||||
return out_mask
|
||||
|
||||
def build_chunks(self, user_offsets, device):
|
||||
"""把拼接序列按用户边界切成每块 ~chunk_users 个用户，返回 [(s0,s1,mask), ...]。
|
||||
每块块内因果，注意力 O(块内S²) 远小于 O(总S²)。仅 1 次同步(读切分边界)。"""
|
||||
chunk_users = int(CONFIG.get("chunk_users", 16))
|
||||
B = user_offsets.numel() - 1 # 用户数（读 shape，无同步）
|
||||
idx = list(range(0, B + 1, chunk_users))
|
||||
if idx[-1] != B:
|
||||
idx.append(B)
|
||||
bounds = user_offsets[idx].tolist() # 1 次同步：取各块的 token 边界
|
||||
chunks = []
|
||||
for c in range(len(bounds) - 1):
|
||||
s0, s1 = bounds[c], bounds[c + 1]
|
||||
local_off = user_offsets[idx[c]:idx[c + 1] + 1] - s0 # 该块内的用户边界（GPU）
|
||||
m = self.causal_mask_syncfree(local_off, s1 - s0, device).unsqueeze(0).unsqueeze(0)
|
||||
chunks.append((s0, s1, m))
|
||||
return chunks
|
||||
|
||||
def causal_mask_syncfree(self, user_offsets, S, device):
|
||||
"""与 get_sequence_causal_mask 等价，但用 searchsorted 求每个位置的用户号，
|
||||
避免 repeat_interleave(张量repeats) 的隐式同步。"""
|
||||
@@ -982,21 +624,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):
|
||||
if batch.get("rep") is not None:
|
||||
seq_input = batch["rep"] # collate 已算好（不计时），跳过 embedding 层
|
||||
elif self._rep_cache is not None:
|
||||
seq_input = self._gather_rep(batch) # load_model 预计算缓存
|
||||
else:
|
||||
seq_input = self.rep_encoder(batch)
|
||||
seq_input = self.rep_encoder(batch)
|
||||
user_offsets = batch["user_offsets"]
|
||||
attn = _resolve_attn(seq_input.device)
|
||||
if attn == "triton":
|
||||
meta = _triton_block_meta(user_offsets, CONFIG.get("triton_block_m", 64),
|
||||
seq_input.device, seq_input.shape[0])
|
||||
extension = {"triton_meta": meta}
|
||||
elif attn == "chunked":
|
||||
extension = {"chunks": self.build_chunks(user_offsets, seq_input.device)}
|
||||
elif attn == "varlen":
|
||||
if attn == "varlen":
|
||||
extension = {"varlen_offsets": user_offsets}
|
||||
elif attn == "flex":
|
||||
S = seq_input.shape[0] # rep_encoder 输出 [S, D]，S=总 token 数
|
||||
@@ -1011,96 +642,10 @@ class CTRModel(nn.Module):
|
||||
encoder_output, moe_loss = self.seq_encoder(x=seq_input, extension=extension)
|
||||
encoder_output = encoder_output.squeeze(0)
|
||||
pred = self.linear(encoder_output)
|
||||
bias = CONFIG.get("logit_bias", 0.0)
|
||||
if bias != 0.0:
|
||||
pred = pred + bias # PCOC 校准（单调，不改 AUC）
|
||||
pred_logits = torch.clamp(pred, min=-15.0, max=15.0)
|
||||
return pred_logits, moe_loss
|
||||
|
||||
|
||||
# ============================================================
|
||||
# RepEncoder 预计算缓存
|
||||
# ============================================================
|
||||
|
||||
def _load_test_user_items(ds_dir):
|
||||
"""流式只加载"测试用户"的 item（避免全量 OOM）。返回 item_dict（仅测试用户）。"""
|
||||
test_csv = ds_dir / "test.csv"
|
||||
history = ds_dir / "history"
|
||||
test_users = set()
|
||||
with open(test_csv) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split(",")
|
||||
if len(parts) >= 2:
|
||||
test_users.add(int(parts[1]))
|
||||
files = (sorted(history.glob("*.csv")) if history.exists() else []) + [test_csv]
|
||||
item_dict = {}
|
||||
for fp in files:
|
||||
has_clk = _detect_has_clk(fp)
|
||||
min_parts = 5 if has_clk else 4
|
||||
with open(fp) as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split(",")
|
||||
if len(parts) < min_parts:
|
||||
continue
|
||||
if int(parts[1]) not in test_users:
|
||||
continue
|
||||
logid = int(parts[0])
|
||||
fs = 5 if has_clk else 4
|
||||
signs, slots = [], []
|
||||
for pair in parts[fs:]:
|
||||
if ":" in pair:
|
||||
s, sl = pair.split(":", 1)
|
||||
signs.append(int(s))
|
||||
slots.append(int(sl))
|
||||
item_dict[logid] = {
|
||||
"signs": np.array(signs, dtype=np.int64),
|
||||
"slots": np.array(slots, dtype=np.int64),
|
||||
}
|
||||
return item_dict
|
||||
|
||||
|
||||
def build_rep_cache(model, item_dict, max_feasign_per_slot, device, chunk=4000, max_slot_id=28):
|
||||
"""直接从 item_dict 逐 item 预计算 RepEncoder 向量（不建 CTRTestSeqDataset，省内存）。
|
||||
|
||||
每个 item 作为一个 segment，逐 slot 拼 values/offsets，跑 model.rep_encoder，
|
||||
与 model(batch) 内的 RepEncoder 输出逐位一致。必须用与评测端一致的
|
||||
max_feasign_per_slot（基线 {1:2}），否则缓存向量与 batch 实际特征不符。
|
||||
"""
|
||||
logids_sorted = sorted(item_dict.keys())
|
||||
emb_chunks = []
|
||||
model.eval()
|
||||
with torch.inference_mode():
|
||||
for i in range(0, len(logids_sorted), chunk):
|
||||
cl = logids_sorted[i:i + chunk]
|
||||
slot_vals = {s: [] for s in range(1, max_slot_id + 1)}
|
||||
slot_offs = {s: [0] for s in range(1, max_slot_id + 1)}
|
||||
for lid in cl:
|
||||
rec = item_dict[lid]
|
||||
by = defaultdict(list)
|
||||
for s, sl in zip(rec["signs"].tolist(), rec["slots"].tolist()):
|
||||
by[sl].append(s)
|
||||
for slot in range(1, max_slot_id + 1):
|
||||
ss = by.get(slot, [])
|
||||
if max_feasign_per_slot and max_feasign_per_slot.get(slot, -1) != -1:
|
||||
ss = ss[:max_feasign_per_slot[slot]]
|
||||
slot_vals[slot].extend(ss)
|
||||
slot_offs[slot].append(len(slot_vals[slot]))
|
||||
batch = {slot: (torch.tensor(slot_vals[slot], dtype=torch.long, device=device),
|
||||
torch.tensor(slot_offs[slot], dtype=torch.long, device=device))
|
||||
for slot in range(1, max_slot_id + 1)}
|
||||
emb_chunks.append(model.rep_encoder(batch)) # [len(cl), d_model]
|
||||
logids = torch.tensor(logids_sorted, dtype=torch.long, device=device) # 已有序
|
||||
emb = torch.cat(emb_chunks)
|
||||
model._rep_cache = (logids.contiguous(), emb.contiguous())
|
||||
return model._rep_cache
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 模型加载入口
|
||||
# ============================================================
|
||||
@@ -1155,16 +700,18 @@ def load_model(ckpt_path, device='cuda:0'):
|
||||
|
||||
if CONFIG["fp16"]:
|
||||
model = model.half()
|
||||
# 默认 Embedding 保留 FP32；emb_fp16=True 时保持 FP16（查表带宽减半）
|
||||
# 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)
|
||||
emb_note = "emb=FP16" if CONFIG.get("emb_fp16", False) else "emb=FP32"
|
||||
print(f"[INFO] FP16 on; {emb_note}; extra FP32-kept: "
|
||||
f"{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)")
|
||||
@@ -1183,38 +730,6 @@ 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 缓存（不计时阶段）===
|
||||
# 优先用"捕获的评测端 item_dict"（不猜路径、不重载、max_feasign 必一致、gather 必命中）；
|
||||
# 捕获不到才退而流式加载 dataset/。任何异常都回退 in-batch RepEncoder。
|
||||
if CONFIG.get("precompute_rep", False) and model._rep_cache is None:
|
||||
try:
|
||||
item_dict = _CAPTURED.get("item_dict")
|
||||
mf = _CAPTURED.get("max_feasign") or {1: 2}
|
||||
source = "captured"
|
||||
if item_dict is None: # 没捕获到 → 退而流式加载 dataset/
|
||||
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:
|
||||
item_dict = _load_test_user_items(ds_dir)
|
||||
source = "stream-loaded"
|
||||
if item_dict is not None:
|
||||
keep = _CAPTURED.get("keep_users")
|
||||
if keep is not None and source == "captured": # 捕获的全量 item_dict → 过滤到测试用户
|
||||
item_dict = {l: r for l, r in item_dict.items()
|
||||
if r.get("userid") in keep}
|
||||
build_rep_cache(model, item_dict, mf, dev)
|
||||
print(f"[INFO] rep cache built ({source}, mf={mf}): "
|
||||
f"{model._rep_cache[0].numel()} items")
|
||||
else:
|
||||
print("[INFO] no data to precompute, fallback to in-batch RepEncoder")
|
||||
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)
|
||||
@@ -1222,22 +737,6 @@ def load_model(ckpt_path, device='cuda:0'):
|
||||
except Exception as e:
|
||||
print(f"[WARNING] torch.compile failed ({e}), running eager")
|
||||
|
||||
global _MODEL_REF
|
||||
_MODEL_REF = model # 供 collate_fn 就地算 RepEncoder
|
||||
|
||||
# 预热 Triton kernel（不计时阶段触发 JIT 编译，避免首个 model(batch) 含编译时间）
|
||||
if _resolve_attn(dev) == "triton":
|
||||
try:
|
||||
H, Dh = model.seq_encoder.n_heads, model.seq_encoder.head_dim
|
||||
dummy_off = torch.tensor([0, 64, 130], device=dev)
|
||||
dq = torch.randn(1, H, 130, Dh, device=dev, dtype=torch.float16)
|
||||
meta = _triton_block_meta(dummy_off, CONFIG.get("triton_block_m", 64), dev, 130)
|
||||
_triton_varlen_attn(dq, dq, dq, meta)
|
||||
torch.cuda.synchronize()
|
||||
print("[INFO] triton kernel warmed up")
|
||||
except Exception as e:
|
||||
print(f"[WARNING] triton warmup failed ({e})")
|
||||
|
||||
print(f"[INFO] Model ready. Device: {dev}")
|
||||
return model, dev
|
||||
|
||||
|
||||
@@ -64,130 +64,6 @@ def test_moe_dense_matches_loop():
|
||||
print(f"[PASS] MoE 稠密向量化 == 逐expert循环 (max err={err:.2e}, dev={dev})")
|
||||
|
||||
|
||||
def test_chunked_matches_dense_attention():
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
rep = infer.RepEncoder(vocab_size=100, emb_dim=8, slot_num=28, d_model=8)
|
||||
seq = infer.TransformerEncoder(d_model=8, n_heads=2, num_layers=1, dim_ff=16)
|
||||
model = infer.CTRModel(rep, seq, d_model=8).to(dev)
|
||||
torch.manual_seed(0)
|
||||
H, Dh = 8, 64
|
||||
offs = _offsets([10, 25, 7, 40, 18, 5, 33], dev) # 7 个用户
|
||||
S = int(offs[-1])
|
||||
q = torch.randn(1, H, S, Dh, device=dev)
|
||||
k = torch.randn(1, H, S, Dh, device=dev)
|
||||
v = torch.randn(1, H, S, Dh, device=dev)
|
||||
with torch.no_grad():
|
||||
dense = infer.scaled_dot_product(q, k, v, {"mask": _dense_causal_mask(offs)[None, None]})
|
||||
infer.CONFIG["chunk_users"] = 3 # 每块 3 个用户
|
||||
chunks = model.build_chunks(offs, torch.device(dev))
|
||||
chunked = infer.scaled_dot_product(q, k, v, {"chunks": chunks})
|
||||
err = (dense - chunked).abs().max().item()
|
||||
assert torch.allclose(dense, chunked, atol=1e-4, rtol=1e-4), f"chunked 不等价 max err={err:.3e}"
|
||||
print(f"[PASS] chunked SDPA == 稠密SDPA (max err={err:.2e}, dev={dev})")
|
||||
|
||||
|
||||
def test_collate_dedup_matches():
|
||||
import numpy as _np
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
enc = infer.RepEncoder(vocab_size=200, emb_dim=512, slot_num=28, d_model=512).to(dev).eval()
|
||||
N = 5
|
||||
plain, dedup = {}, {}
|
||||
for s in range(1, 29):
|
||||
seg_vals, offs_p = [], [0]
|
||||
u_vals, u_w, offs_d = [], [], [0]
|
||||
for _ in range(N):
|
||||
m = int(torch.randint(1, 8, (1,)))
|
||||
signs = torch.randint(0, 200, (m,)).tolist()
|
||||
signs = signs + signs[:max(0, m - 1)] # 制造段内重复
|
||||
seg_vals.extend(signs); offs_p.append(len(seg_vals))
|
||||
uq, ct = _np.unique(_np.asarray(signs), return_counts=True)
|
||||
u_vals.extend(uq.tolist()); u_w.extend(ct.tolist()); offs_d.append(len(u_vals))
|
||||
plain[s] = (torch.tensor(seg_vals, device=dev), torch.tensor(offs_p, device=dev))
|
||||
dedup[s] = (torch.tensor(u_vals, device=dev), torch.tensor(offs_d, device=dev),
|
||||
torch.tensor(u_w, dtype=torch.float32, device=dev))
|
||||
with torch.no_grad():
|
||||
infer.CONFIG["use_embedding_bag"] = True
|
||||
ref = enc(plain)
|
||||
new = enc(dedup)
|
||||
infer.CONFIG["use_embedding_bag"] = False
|
||||
err = (ref - new).abs().max().item()
|
||||
assert torch.allclose(ref, new, atol=1e-3, rtol=1e-3), f"collate_dedup 不等价 max err={err:.3e}"
|
||||
print(f"[PASS] collate_dedup(去重+计数) == 全展开 (max err={err:.2e}, dev={dev})")
|
||||
|
||||
|
||||
def test_embedding_bag_matches():
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
slot_num, emb_dim, d_model = 28, 512, 512
|
||||
enc = infer.RepEncoder(vocab_size=200, emb_dim=emb_dim, slot_num=slot_num,
|
||||
d_model=d_model).to(dev).eval()
|
||||
N = 6
|
||||
batch = {}
|
||||
for s in range(1, slot_num + 1):
|
||||
counts = torch.randint(0, 8, (N,))
|
||||
vals = torch.randint(0, 200, (int(counts.sum()),), device=dev)
|
||||
offs = torch.cat([torch.zeros(1, dtype=torch.long), counts.cumsum(0)]).to(dev)
|
||||
batch[s] = (vals, offs)
|
||||
with torch.no_grad():
|
||||
infer.CONFIG["use_embedding_bag"] = False
|
||||
ref = enc(batch)
|
||||
infer.CONFIG["use_embedding_bag"] = True
|
||||
new = enc(batch)
|
||||
infer.CONFIG["use_embedding_bag"] = False
|
||||
err = (ref - new).abs().max().item()
|
||||
assert torch.allclose(ref, new, atol=1e-3, rtol=1e-3), f"embedding_bag 不等价 max err={err:.3e}"
|
||||
print(f"[PASS] embedding_bag == segment_reduce (max err={err:.2e}, dev={dev})")
|
||||
|
||||
|
||||
def test_sparse_pool_matches():
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
slot_num, emb_dim, d_model = 28, 512, 512
|
||||
enc = infer.RepEncoder(vocab_size=200, emb_dim=emb_dim, slot_num=slot_num,
|
||||
d_model=d_model).to(dev).eval()
|
||||
N = 6
|
||||
batch = {}
|
||||
for s in range(1, slot_num + 1):
|
||||
counts = torch.randint(0, 8, (N,))
|
||||
# 故意制造段内重复：值域很小，重复率高
|
||||
vals = torch.randint(0, 30, (int(counts.sum()),), device=dev)
|
||||
offs = torch.cat([torch.zeros(1, dtype=torch.long), counts.cumsum(0)]).to(dev)
|
||||
batch[s] = (vals, offs)
|
||||
with torch.no_grad():
|
||||
infer.CONFIG["sparse_pool"] = False
|
||||
infer.CONFIG["dedup_embedding"] = True
|
||||
ref = enc(batch)
|
||||
infer.CONFIG["sparse_pool"] = True
|
||||
new = enc(batch)
|
||||
infer.CONFIG["sparse_pool"] = False
|
||||
err = (ref - new).abs().max().item()
|
||||
assert torch.allclose(ref, new, atol=2e-2, rtol=2e-2), f"sparse_pool 不等价 max err={err:.3e}"
|
||||
print(f"[PASS] sparse_pool == segment_reduce (max err={err:.2e}, dev={dev})")
|
||||
|
||||
|
||||
def test_triton_varlen_matches_dense():
|
||||
if not (torch.cuda.is_available() and infer._HAS_TRITON):
|
||||
print("[SKIP] Triton varlen 等价测试（需 CUDA + triton）")
|
||||
return
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
H, Dh = 8, 64
|
||||
offs = _offsets([10, 64, 1, 130, 64, 200], dev) # 含跨多块/单token/正好整块的段
|
||||
S = int(offs[-1])
|
||||
q = torch.randn(1, H, S, Dh, device=dev, dtype=torch.float16)
|
||||
k = torch.randn(1, H, S, Dh, device=dev, dtype=torch.float16)
|
||||
v = torch.randn(1, H, S, Dh, device=dev, dtype=torch.float16)
|
||||
with torch.no_grad():
|
||||
dense = infer.scaled_dot_product(q, k, v, {"mask": _dense_causal_mask(offs)[None, None]})
|
||||
meta = infer._triton_block_meta(offs, 64, q.device, S)
|
||||
trit = infer.scaled_dot_product(q, k, v, {"triton_meta": meta})
|
||||
err = (dense.float() - trit.float()).abs().max().item()
|
||||
assert torch.allclose(dense.float(), trit.float(), atol=3e-2, rtol=3e-2), \
|
||||
f"Triton varlen 不等价 max err={err:.3e}"
|
||||
print(f"[PASS] Triton varlen flash == 稠密SDPA (max err={err:.2e})")
|
||||
|
||||
|
||||
def test_syncfree_mask_matches():
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
rep = infer.RepEncoder(vocab_size=100, emb_dim=8, slot_num=28, d_model=8)
|
||||
@@ -222,25 +98,6 @@ def test_varlen_matches_dense_attention():
|
||||
print(f"[PASS] varlen(嵌套张量) == 稠密SDPA (max err={err:.2e})")
|
||||
|
||||
|
||||
def test_sparse_moe_matches_dense():
|
||||
# 大 capacity（无丢弃）下，稀疏 MoE 应与 dense 数学等价
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
m = infer.SMoE(d_model=512, dim_ff=1024, num_experts=8, k=2).to(dev).eval()
|
||||
x = torch.randn(1, 200, 512, device=dev)
|
||||
with torch.no_grad():
|
||||
infer.CONFIG["moe_sparse"] = False
|
||||
ref, _ = m(x)
|
||||
infer.CONFIG["moe_sparse"] = True
|
||||
infer.CONFIG["moe_capacity"] = 8.0 # 足够大，不丢 token
|
||||
new, _ = m(x)
|
||||
infer.CONFIG["moe_sparse"] = False
|
||||
infer.CONFIG["moe_capacity"] = 1.25
|
||||
err = (ref - new).abs().max().item()
|
||||
assert torch.allclose(ref, new, atol=1e-3, rtol=1e-3), f"sparse MoE 不等价 max err={err:.3e}"
|
||||
print(f"[PASS] sparse MoE(大capacity) == dense (max err={err:.2e}, dev={dev})")
|
||||
|
||||
|
||||
def test_fused_embedding_matches_perslot():
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
@@ -289,14 +146,8 @@ def test_flex_matches_dense_attention():
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_moe_dense_matches_loop()
|
||||
test_sparse_moe_matches_dense()
|
||||
test_fused_embedding_matches_perslot()
|
||||
test_embedding_bag_matches()
|
||||
test_collate_dedup_matches()
|
||||
test_sparse_pool_matches()
|
||||
test_syncfree_mask_matches()
|
||||
test_triton_varlen_matches_dense()
|
||||
test_chunked_matches_dense_attention()
|
||||
test_varlen_matches_dense_attention()
|
||||
test_flex_matches_dense_attention()
|
||||
print("[DONE] 等价测试结束")
|
||||
|
||||
@@ -1,624 +0,0 @@
|
||||
# CTI 生成式推荐广告排序推理优化方案
|
||||
|
||||
> 基于 baseline 代码分析、HSTU / GRAB 论文研究、官方提交规范
|
||||
|
||||
---
|
||||
|
||||
## 0. 前置修复：接口对齐（必须首先完成）
|
||||
|
||||
### 0.1 发现的问题
|
||||
|
||||
官方 `任务提交接口说明.md` 定义了评测系统调用接口。baseline 代码与规范存在三处致命不匹配：
|
||||
|
||||
| 接口 | 官方要求 | Baseline 实际 | 后果 |
|
||||
|------|----------|---------------|------|
|
||||
| 数据集类名 | `CTRTestSeqDataset` | `CTRUserDataset` | 评测 `from infer import CTRTestSeqDataset` 失败 |
|
||||
| 构造参数 | `test_logids_ordered`, `item_dict`, `user_seq`, `max_feasign_per_slot`, `max_ctx_len` | `item_dict`, `user_seq`, `max_feasign_per_slot`, `pred_logids` | 参数名和数量不匹配 |
|
||||
| `load_model` | `load_model(ckpt_path: Path) -> (model, device)` | `load_model(device='cuda:0', ckpt_path=None)` | 评测调用时 Path 会被错误赋给 `device` 参数 |
|
||||
|
||||
评测系统会用 `python -c "from infer import CTRTestSeqDataset, load_model; ..."` 来加载你的代码。**这三处不改，任何优化都白费。**
|
||||
|
||||
### 0.2 修复步骤
|
||||
|
||||
**修复 1：重命名类并调整构造参数**
|
||||
|
||||
将 `CTRUserDataset` 改为 `CTRTestSeqDataset`，参数名改为 `test_logids_ordered`，增加 `max_ctx_len` 占位：
|
||||
|
||||
```python
|
||||
class CTRTestSeqDataset(Dataset):
|
||||
"""按用户组织的 CTR 测试数据集（对齐评测接口）"""
|
||||
|
||||
def __init__(self, test_logids_ordered, item_dict, user_seq=None,
|
||||
max_feasign_per_slot=None, max_ctx_len=None):
|
||||
super().__init__()
|
||||
self.item_dict = item_dict
|
||||
self.user_seq = user_seq if user_seq else {}
|
||||
self.max_feasign_per_slot = max_feasign_per_slot
|
||||
self.max_ctx_len = max_ctx_len
|
||||
self.pred_logids = set(test_logids_ordered) if test_logids_ordered else set()
|
||||
# ... 其余逻辑不变
|
||||
```
|
||||
|
||||
**修复 2：修正 `load_model` 签名**
|
||||
|
||||
```python
|
||||
def load_model(ckpt_path, device='cuda:0'):
|
||||
"""加载模型。签名对齐评测接口：第一个参数必须是 ckpt_path。"""
|
||||
# ... 其余逻辑不变
|
||||
```
|
||||
|
||||
### 0.3 提交验证标准
|
||||
|
||||
修改后必须在 AI Studio 提交一次，确认能跑通（得分 > 0），再开始做优化。这是所有优化的前提。
|
||||
|
||||
---
|
||||
|
||||
## 1. 优化起点（Baseline 数据）
|
||||
|
||||
| 指标 | 当前值 | 说明 |
|
||||
|------|--------|------|
|
||||
| 推理耗时 | 229.18s | **只计 `model(batch)` 的逐 batch 累加时间** |
|
||||
| AUC | 0.759 | 阈值 ≥ 0.65 |
|
||||
| PCOC | 1.110 | 阈值 [0.85, 1.15] |
|
||||
| 综合得分 | 25.85 | `score_latency * 70 + score_model * 30` |
|
||||
|
||||
技术栈：
|
||||
- Python 3.10.10 + PyTorch 2.6.0 + CUDA 12.4
|
||||
- 模型：RepEncoder（28 slot / 512 维 embedding）→ 8 层 Transformer（8 头 / 512 维）→ MoE FFN（8 个 expert，Top-2 gating）
|
||||
- 数据：已缓存为 9 个 shard 分片（shard_0000.pt ~ shard_0008.pt），共 2039 batch / 7774 条预测
|
||||
|
||||
---
|
||||
|
||||
## 2. 约束条件（全部来自官方规则）
|
||||
|
||||
| 约束 | 来源 | 说明 |
|
||||
|------|------|------|
|
||||
| 不能改"组网" | 赛题说明 | 模型结构、层数、头数、维度不可改 |
|
||||
| 不能改参数（权重值） | 赛题说明 | 量化/稀疏化/剪枝**明确允许** |
|
||||
| 不能用测试集训练 | 赛题说明 | 仅推理，不做任何训练 |
|
||||
| 推理时限 ≤ 300s | 提交规范 §4.1 | 超时总分直接 0。**只计 `model(batch)` 时间，数据加载和模型加载不计** |
|
||||
| build_env.sh 时限 ≤ 720s | 提交规范 §3 | 超时或非 0 退出码直接失败 |
|
||||
| AUC ≥ 0.65 且 PCOC ∈ [0.85, 1.15] | 提交规范 §4.2 | 任一不满足，总分直接 0 |
|
||||
| 压缩包内不能有 `dataset/` 和 `ckpt.pt` | 提交规范 §2.2 | 评测系统自行提供 |
|
||||
| 压缩包后缀必须是 `.zip`/`.tar.gz`/`.tar` | 提交规范 §1 | 其他格式不识别 |
|
||||
| 解压后文件必须直接在根目录 | 提交规范 §1 | 不能多一层包裹文件夹 |
|
||||
|
||||
---
|
||||
|
||||
## 3. 优化方案总览（修订版）
|
||||
|
||||
```
|
||||
第〇步：接口对齐（不优化，先跑通）
|
||||
└── 确认能提交得分 > 0
|
||||
│
|
||||
第①步：FP16 量化 ←── 明确允许，收益最大
|
||||
└── 预期：229s → ~120s
|
||||
│
|
||||
第②步：Flash Attention ←── 数学等价，不改组网
|
||||
└── 预期：120s → ~90s
|
||||
│
|
||||
第③步：torch.compile ←── 编译器优化，不改组网
|
||||
└── 预期：90s → ~65s
|
||||
│
|
||||
第④步：数据流优化 ←── 减少 CPU→GPU 传输开销
|
||||
└── 预期：65s → ~55s
|
||||
│
|
||||
第⑤步：MoE 优化 ←── "剪枝"明确允许
|
||||
└── 预期：55s → ~50s
|
||||
│
|
||||
第⑥步：INT8 量化（可选） ←── 收益大但风险高
|
||||
└── 只有在①②③⑤之后仍不够时才尝试
|
||||
```
|
||||
|
||||
每个方案完成后在 AI Studio 提交验证，确认分数提升后再进入下一步。
|
||||
|
||||
---
|
||||
|
||||
## 4. 方案一：FP16 量化
|
||||
|
||||
### 4.1 合规性
|
||||
|
||||
✅ 比赛规则明确写有"量化除外"，属于允许范围。
|
||||
|
||||
### 4.2 原理
|
||||
|
||||
当前 baseline 使用 FP32。FP16 将模型参数和激活值减半，GPU 的张量核心对 FP16 有原生加速。注意 `Embedding` 层保留 FP32（索引是整数，embedding 查表时再转换）。
|
||||
|
||||
### 4.3 实现
|
||||
|
||||
修改 `load_model` 函数中的模型加载部分：
|
||||
|
||||
```python
|
||||
def load_model(ckpt_path, device='cuda:0'):
|
||||
# ... 模型初始化不变 ...
|
||||
|
||||
dev = torch.device(device if torch.cuda.is_available() else "cpu")
|
||||
|
||||
if ckpt_path is None:
|
||||
ckpt_path = Path(__file__).parent / 'ckpt.pt'
|
||||
else:
|
||||
ckpt_path = Path(ckpt_path)
|
||||
|
||||
if ckpt_path.exists():
|
||||
ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
|
||||
model.load_state_dict(ckpt['model_state_dict'])
|
||||
|
||||
# === FP16 优化：将模型转为半精度 ===
|
||||
model = model.half()
|
||||
# Embedding 层保留 FP32（索引是 int，不需要转）
|
||||
model.rep_encoder.emb = model.rep_encoder.emb.to(torch.float32)
|
||||
|
||||
print(f"[INFO] Loaded checkpoint from {ckpt_path} (epoch={ckpt.get('epoch', '?')})")
|
||||
else:
|
||||
print(f"[WARNING] Checkpoint {ckpt_path} not found, using random weights")
|
||||
|
||||
model.to(dev)
|
||||
model.eval()
|
||||
print(f"[INFO] Model ready. Device: {dev}")
|
||||
return model, dev
|
||||
```
|
||||
|
||||
`move_batch_to_device` 中同步处理输入数据精度：
|
||||
|
||||
```python
|
||||
def move_batch_to_device(batch, device):
|
||||
if isinstance(batch, dict):
|
||||
return {k: move_batch_to_device(v, device) for k, v in batch.items()}
|
||||
elif isinstance(batch, (list, tuple)):
|
||||
return [move_batch_to_device(x, device) for x in batch]
|
||||
elif torch.is_tensor(batch):
|
||||
x = batch.to(device)
|
||||
# 浮点 tensor → FP16，整数 tensor 保持不变
|
||||
if x.dtype == torch.float32:
|
||||
x = x.half()
|
||||
return x
|
||||
else:
|
||||
return batch
|
||||
```
|
||||
|
||||
同时修改 `main()` 中缓存 batch 的逻辑，在落地磁盘前就转 FP16，避免推理时逐 batch 转换：
|
||||
|
||||
```python
|
||||
for batch in test_loader:
|
||||
# 预转为 FP16
|
||||
batch = move_batch_to_device(batch, torch.device('cpu'))
|
||||
all_batches.append(batch)
|
||||
```
|
||||
|
||||
### 4.4 风险
|
||||
|
||||
- 极低。现代 NVIDIA GPU（V100/A100/H100 等）的 FP16 吞吐量是 FP32 的 2-8 倍
|
||||
- Embedding 保留 FP32 是因为 embedding 查表操作用 int 索引，不受浮点精度影响
|
||||
|
||||
### 4.5 预期收益
|
||||
|
||||
| 指标 | 变化 |
|
||||
|------|------|
|
||||
| 推理时间 | 229s → ~115-150s |
|
||||
| AUC | 几乎不变（差异 < 0.0001） |
|
||||
|
||||
---
|
||||
|
||||
## 5. 方案二：Flash Attention
|
||||
|
||||
### 5.1 合规性
|
||||
|
||||
✅ Flash Attention 是**数学等价**的注意力计算。它用分块（tiling）算法计算 `softmax(QK^T/√d)V`，**输出结果与标准 attention 相同**（在浮点误差范围内）。不改变模型结构、不改变参数——只改变了内存访问模式和计算顺序。这属于编译器级优化，不是组网修改。
|
||||
|
||||
### 5.2 原理
|
||||
|
||||
Baseline 的 `scaled_dot_product` 先计算完整的 `QK^T` 矩阵（O(L²) 显存），再做 softmax，再乘 V。Flash Attention 将 Q、K、V 分块加载到 SRAM，逐块计算 softmax 并累积结果，**避免完整 QK^T 驻留 HBM**。
|
||||
|
||||
Baseline 用了 Sequence Packing（同一用户的多个 impression 拼成一条长序列），序列越长收益越大。
|
||||
|
||||
### 5.3 实现
|
||||
|
||||
PyTorch 2.0+ 提供了 `F.scaled_dot_product_attention`，自动选择最优后端（Flash Attention / Memory Efficient Attention / 标准实现）：
|
||||
|
||||
```python
|
||||
import torch.nn.functional as F
|
||||
|
||||
def scaled_dot_product(q, k, v, extension):
|
||||
"""使用 PyTorch SDPA 后端（自动启用 Flash Attention）"""
|
||||
d = q.size(-1)
|
||||
|
||||
if extension is not None and "mask" in extension:
|
||||
mask = extension["mask"]
|
||||
# mask 形状: [1, 1, S, S] 或 [B, 1, S, S]
|
||||
# 转换为 float mask，确保 device 和 dtype 一致
|
||||
attn_mask = mask.to(device=q.device, dtype=q.dtype)
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
return F.scaled_dot_product_attention(
|
||||
q, k, v,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
)
|
||||
```
|
||||
|
||||
验证是否启用了 Flash Attention：
|
||||
|
||||
```python
|
||||
import torch
|
||||
# 在模型加载后运行，确认后端可用
|
||||
print("Flash SDP:", torch.backends.cuda.flash_sdp_enabled())
|
||||
print("Mem Efficient SDP:", torch.backends.cuda.mem_efficient_sdp_enabled())
|
||||
print("Math SDP:", torch.backends.cuda.math_sdp_enabled())
|
||||
```
|
||||
|
||||
### 5.4 与 Baseline 输出的兼容性
|
||||
|
||||
Baseline 中 mask 是 `torch.bool` 类型，原代码用 `scores.masked_fill(mask == 0, float("-inf"))` 处理。`F.scaled_dot_product_attention` 接受 bool mask 时自动做等价处理。数学等价，不会影响 AUC/PCOC。
|
||||
|
||||
### 5.5 预期收益
|
||||
|
||||
| 序列长度 | 加速比 |
|
||||
|----------|--------|
|
||||
| < 512 | 1.1x - 1.3x |
|
||||
| 512 - 2048 | 1.5x - 2x |
|
||||
| > 2048 | 2x - 4x |
|
||||
|
||||
---
|
||||
|
||||
## 6. 方案三：torch.compile
|
||||
|
||||
### 6.1 合规性
|
||||
|
||||
✅ `torch.compile` 是纯粹的**编译器优化**（JIT 编译 + 算子融合）。它不改变模型结构、不改变权重值、不改变任何数学运算——只是把多个小 kernel 合并成一个大 kernel 以减少 GPU kernel launch 开销。
|
||||
|
||||
### 6.2 实现
|
||||
|
||||
在 `load_model` 中，`model.eval()` 之后添加一行：
|
||||
|
||||
```python
|
||||
model = torch.compile(model, mode="reduce-overhead")
|
||||
```
|
||||
|
||||
`mode` 选择：
|
||||
| 模式 | 行为 | 适合场景 |
|
||||
|------|------|----------|
|
||||
| `"default"` | 平衡编译速度与运行时 | 不确定时先用这个 |
|
||||
| `"reduce-overhead"` | 更激进的融合，减少 kernel launch | **推理场景首选** |
|
||||
| `"max-autotune"` | 自动调优 Triton kernel（首次编译慢） | 最终提交时切换到此 |
|
||||
|
||||
建议先用 `"reduce-overhead"` 快速验证，最终提交时换成 `"max-autotune"`。
|
||||
|
||||
### 6.3 已知限制
|
||||
|
||||
- 首次运行会触发 JIT 编译（0.5-2 分钟），在评测环境可能计入推理时间
|
||||
- 解决：在 `build_env.sh` 中用一个小输入做一次预热编译：
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
set -e
|
||||
pip install -r requirements.txt -i https://mirrors.aliyun.com/pypi/simple/
|
||||
|
||||
# 预热 torch inductor，避免推理时编译
|
||||
python -c "
|
||||
import torch
|
||||
@torch.compile(mode='max-autotune')
|
||||
def warmup(x):
|
||||
return x * 2
|
||||
x = torch.randn(100, 100, device='cuda')
|
||||
warmup(x)
|
||||
print('Inductor cache ready')
|
||||
"
|
||||
|
||||
echo "build env success"
|
||||
```
|
||||
|
||||
### 6.4 预期收益
|
||||
|
||||
| 模式 | 加速比 |
|
||||
|------|--------|
|
||||
| reduce-overhead | 1.2x - 1.4x |
|
||||
| max-autotune | 1.3x - 1.5x（首次多花 1-2 分钟） |
|
||||
|
||||
---
|
||||
|
||||
## 7. 方案四：数据流优化
|
||||
|
||||
### 7.1 重要澄清：评测的计时逻辑
|
||||
|
||||
从提交规范 §4.1：
|
||||
> 只统计模型前向时间（逐 batch 累加 `model(batch)` 耗时），数据加载、模型加载不计入。
|
||||
|
||||
所以 **DataLoader 层面的优化（num_workers、pin_memory）不影响评分**。但这不代表数据流优化没用——以下两点仍然有效：
|
||||
|
||||
1. **减少 `move_batch_to_device` 开销**：此函数在 `model(batch)` 之前调用，但如果它在推理循环里逐 batch 执行，仍会阻塞。解决方案：提前把所有 batch 搬到 GPU。
|
||||
2. **减少 FP32→FP16 转换开销**：在缓存阶段就转好，推理循环里无需重复转换。
|
||||
|
||||
### 7.2 实现
|
||||
|
||||
在 `main()` 的数据加载部分，缓存时就完成设备搬运和类型转换：
|
||||
|
||||
```python
|
||||
# 原代码
|
||||
all_batches = [batch for batch in test_loader]
|
||||
|
||||
# 优化：缓存时直接搬到 GPU + 转 FP16
|
||||
dev = torch.device('cuda:0')
|
||||
all_batches = []
|
||||
for batch in test_loader:
|
||||
batch = move_batch_to_device(batch, dev)
|
||||
all_batches.append(batch)
|
||||
```
|
||||
|
||||
这样推理循环中不需要再调用 `move_batch_to_device`：
|
||||
|
||||
```python
|
||||
with torch.no_grad():
|
||||
for batch in tqdm(all_batches, desc="Inference"):
|
||||
# batch 已在 GPU 上，无需 move_batch_to_device
|
||||
pred_mask = batch["pred_mask"].bool()
|
||||
t_start = time.time()
|
||||
logits, moe_loss = model(batch)
|
||||
logits = logits.squeeze(-1)
|
||||
probs = torch.sigmoid(logits)
|
||||
time_sum += time.time() - t_start
|
||||
# ...
|
||||
```
|
||||
|
||||
### 7.3 潜在问题
|
||||
|
||||
全部 batch 都在 GPU 上会占用大量显存。如果数据总量超过可用显存，会 OOM。安全做法：先测试单个 shard 的显存占用，若不够则分批预加载。
|
||||
|
||||
如果显存不够，采用双缓冲策略（当前 batch 在 GPU 推理，下一个 batch 异步上传）。
|
||||
|
||||
### 7.4 预期收益
|
||||
|
||||
有限（因为数据加载本来就不计分），但能消除循环内的不必要开销。
|
||||
|
||||
---
|
||||
|
||||
## 8. 方案五：MoE 推理优化
|
||||
|
||||
### 8.1 合规性
|
||||
|
||||
✅ "剪枝"是规则明确允许的三种优化之一（量化/稀疏/剪枝）。Expert 剪枝属于模型剪枝的子类。
|
||||
|
||||
### 8.2 当前状态
|
||||
|
||||
```python
|
||||
class SMoE(nn.Module):
|
||||
def __init__(self, d_model, dim_ff, num_experts=8, k=2):
|
||||
# 每层 8 个 FFN expert，每个 token 激活 2 个
|
||||
```
|
||||
|
||||
8 层 × 8 expert = 64 个 FFN 模块。推理时大多数 token 只会路由到少数 expert。
|
||||
|
||||
### 8.3 优化方向
|
||||
|
||||
**a) Expert 负载统计 + 合并**
|
||||
|
||||
对少量数据做一次前向，统计各 expert 的被激活次数：
|
||||
|
||||
```python
|
||||
expert_hits = torch.zeros(8, 8) # [layers, experts]
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
for batch in sample_batches[:10]: # 只用 10 个 batch 采样
|
||||
batch = move_batch_to_device(batch, dev)
|
||||
for layer_idx, moe in enumerate(model.seq_encoder.moe):
|
||||
topk_idx, _, _ = moe.gate(x) # x 需要从 forward 中间取
|
||||
for e in range(8):
|
||||
expert_hits[layer_idx, e] += (topk_idx == e).sum()
|
||||
```
|
||||
|
||||
如果某个 expert 激活次数 < 1%，可将其权重合并到最相似的 expert（通过权重余弦相似度找最近邻），或直接移除（需验证 AUC 不掉）。
|
||||
|
||||
**b) 替换为静态 FFN（激进方案）**
|
||||
|
||||
如果发现某些层的 expert 负载极度不均衡（例如 95% token 走同一个 expert），可直接把该层的 SMoE 替换为单个 Expert：
|
||||
|
||||
```python
|
||||
# 仅当某层 expert 负载极不均衡时
|
||||
# 找到最常用的 expert
|
||||
best = expert_hits[layer_idx].argmax().item()
|
||||
# 替换 SMoE 为单个 Expert
|
||||
model.seq_encoder.moe[layer_idx] = model.seq_encoder.moe[layer_idx].experts[best]
|
||||
```
|
||||
|
||||
此方案**必须提交验证 AUC/PCOC 不跌破阈值**。
|
||||
|
||||
### 8.4 风险
|
||||
|
||||
- 负载统计需要从 Transformer forward 中间提取 gate 输出，需要修改 forward 添加 hook
|
||||
- Expert 合并可能影响 AUC 0.001-0.005
|
||||
|
||||
### 8.5 预期收益
|
||||
|
||||
| 操作 | 加速比 |
|
||||
|------|--------|
|
||||
| 移除 1-2 个死 expert | 1.05x - 1.15x |
|
||||
| 替换 2-3 层为单 FFN | 1.2x - 1.4x |
|
||||
|
||||
---
|
||||
|
||||
## 9. 方案六：INT8 量化（可选）
|
||||
|
||||
### 9.1 合规性
|
||||
|
||||
✅ "量化除外"，明确允许。
|
||||
|
||||
### 9.2 适用条件
|
||||
|
||||
只有在前 5 步做完后仍需进一步加速时才尝试。INT8 量化需要校准数据（可用测试集的一小部分做 calibration），有精度损失风险。
|
||||
|
||||
```python
|
||||
import torch.quantization as quant
|
||||
|
||||
# 仅对 Transformer encoder 做 INT8 量化（Embedding 层跳过）
|
||||
model.seq_encoder.qconfig = quant.get_default_qconfig('qnnpack')
|
||||
model.seq_encoder = quant.quantize_dynamic(
|
||||
model.seq_encoder,
|
||||
{nn.Linear}, # 仅量化 Linear 层
|
||||
dtype=torch.qint8,
|
||||
)
|
||||
```
|
||||
|
||||
### 9.3 风险
|
||||
|
||||
- AUC 可能下降 0.005-0.02，需提交后验证
|
||||
- 如果评测环境无 GPU（CPU-only），QNNPACK 才有用；GPU 下需用 TensorRT 做 INT8
|
||||
|
||||
---
|
||||
|
||||
## 10. CUDA Graph（为什么不做）
|
||||
|
||||
原方案中列了 CUDA Graph，审查后**移除**，理由如下：
|
||||
|
||||
1. CUDA Graph 要求**所有输入形状完全相同**。Baseline 数据经过 Sequence Packing，不同 batch 的序列长度差异很大，不满足这一前提
|
||||
2. 每个 batch 的 `user_offsets` 长度不同，导致 mask 形状也不同
|
||||
3. 若要强行使用，需要对 batch 做 padding 对齐，反而引入额外开销
|
||||
|
||||
**结论：不适用于本场景，放弃此方案。**
|
||||
|
||||
---
|
||||
|
||||
## 11. 部署要点
|
||||
|
||||
### 11.1 压缩包结构
|
||||
|
||||
```
|
||||
submit.zip
|
||||
├── infer.py # 主推理脚本（实现所有必需接口）
|
||||
├── requirements.txt # Python 依赖列表
|
||||
└── build_env.sh # 环境构建脚本（可选但推荐）
|
||||
```
|
||||
|
||||
提交前务必验证：
|
||||
- [ ] 包内**没有** `dataset/` 目录
|
||||
- [ ] 包内**没有** `ckpt.pt` 文件
|
||||
- [ ] 包内**没有**多余的顶层文件夹
|
||||
- [ ] 后缀是 `.zip`（或 `.tar.gz`）
|
||||
|
||||
### 11.2 build_env.sh 设计
|
||||
|
||||
```bash
|
||||
#!/bin/bash
|
||||
set -e
|
||||
# 安装依赖（评测系统自动使用阿里云镜像）
|
||||
pip install -r requirements.txt
|
||||
|
||||
# 预热 torch compile（如果方案三启用）
|
||||
python -c "
|
||||
import torch
|
||||
@torch.compile(mode='max-autotune')
|
||||
def _warmup(x): return x * 2
|
||||
_warmup(torch.randn(100, 100, device='cuda'))
|
||||
print('Inductor ready')
|
||||
" 2>/dev/null || echo 'torch.compile not available, skipping'
|
||||
|
||||
echo "build env success"
|
||||
```
|
||||
|
||||
### 11.3 requirements.txt（最小化）
|
||||
|
||||
```
|
||||
torch==2.6.0
|
||||
triton==3.2.0
|
||||
numpy==2.2.6
|
||||
scikit-learn==1.7.2
|
||||
tqdm==4.67.3
|
||||
```
|
||||
|
||||
- 去掉所有 `nvidia-*` 包（评测环境已预装 CUDA）
|
||||
- 版本号精确锁定，避免安装时依赖冲突
|
||||
|
||||
---
|
||||
|
||||
## 12. 风险评估与底线策略
|
||||
|
||||
### 各方案对模型质量的影响
|
||||
|
||||
| 方案 | AUC 影响 | PCOC 影响 | 0 分风险 |
|
||||
|------|----------|-----------|----------|
|
||||
| 接口修复 | 无 | 无 | **必须做，否则直接 0 分** |
|
||||
| FP16 | < 0.0001 | < 0.0001 | 极低 |
|
||||
| Flash Attention | < 0.0001 | 无 | 极低 |
|
||||
| torch.compile | 无 | 无 | 低（首次编译可能超时） |
|
||||
| GPU 预加载 | 无 | 无 | 低（OOM 风险） |
|
||||
| MoE 剪枝 | 0.001 - 0.005 | 微小 | 中（需提交验证） |
|
||||
| INT8 | 0.005 - 0.02 | 可能偏移 | 高（可能跌破阈值） |
|
||||
|
||||
### 安全策略（必须遵守）
|
||||
|
||||
每完成一个优化：
|
||||
1. 在 AI Studio 提交，拿到新得分
|
||||
2. 记录本次得分变化
|
||||
3. 如果 AUC < 0.65 或 PCOC 不在 [0.85, 1.15]，立即回退该优化
|
||||
4. 如果得分上升，保留并进入下一步
|
||||
|
||||
### 保底策略
|
||||
|
||||
- 至少完成接口修复 + FP16，能稳定拿到 > 25 分
|
||||
- 如果 Flash Attention / torch.compile 跑不通，回退不影响得分
|
||||
- **在截止日期前一天（6 月 25 日）停止实验**，提交当前最优版本
|
||||
|
||||
---
|
||||
|
||||
## 13. 实施检查清单
|
||||
|
||||
### 第〇周（立即）
|
||||
- [ ] 修复接口：`CTRTestSeqDataset`、`load_model(ckpt_path, ...)`
|
||||
- [ ] 在 AI Studio 提交一次，确认得分 > 0（验证接口正确）
|
||||
|
||||
### 第一轮
|
||||
- [ ] FP16 量化：`model.half()` + embedding 保留 FP32
|
||||
- [ ] 数据预加载到 GPU + 预转 FP16
|
||||
- [ ] 提交验证
|
||||
|
||||
### 第二轮
|
||||
- [ ] Flash Attention：替换 `scaled_dot_product`
|
||||
- [ ] `torch.compile(mode="reduce-overhead")`
|
||||
- [ ] `build_env.sh` 写预热逻辑
|
||||
- [ ] 提交验证
|
||||
|
||||
### 第三轮（时间允许）
|
||||
- [ ] MoE expert 负载分析 + 合并
|
||||
- [ ] `torch.compile` 切换为 `"max-autotune"`
|
||||
- [ ] INT8 量化评估（如果得分仍不满意）
|
||||
|
||||
---
|
||||
|
||||
## 14. 预期效果（修订）
|
||||
|
||||
| 阶段 | 预期推理时间 | 预期得分 | 主要贡献 |
|
||||
|------|-------------|----------|----------|
|
||||
| Baseline | 229s | 25.85 | — |
|
||||
| 接口修复 | 229s | 25.85 | 确保能跑 |
|
||||
| + FP16 + GPU 预加载 | ~120s | ~50 | 量化为主要贡献 |
|
||||
| + Flash Attention | ~90s | ~60 | 长序列受益 |
|
||||
| + torch.compile | ~65s | ~70 | 算子融合 |
|
||||
| + MoE 优化 | ~50s | ~78 | 剪枝 |
|
||||
| 极限（+INT8） | ~30s | ~87 | 有精度风险 |
|
||||
|
||||
---
|
||||
|
||||
## 附录 A：官方规则原文引用
|
||||
|
||||
来自赛题说明页面：
|
||||
> 【由于是推理性能优化，组网不可进行策略性改动，不可对测试集进行训练】
|
||||
|
||||
来自 baseline notebook 第三单元：
|
||||
> 选手不可对组网和相关参数进行修改。【违规为0分】
|
||||
> **量化稀疏剪枝除外**
|
||||
|
||||
来自 `任务提交接口说明.md` §4-5：
|
||||
> - 延迟阈值：300 秒。只统计模型前向时间（逐 batch 累加 model(batch) 耗时）
|
||||
> - AUC ∈ [0.65, 1.0]、PCOC ∈ [0.85, 1.15]，任一不满足总分直接置 0
|
||||
> - 压缩包自带 dataset/ 或 ckpt.pt → 评测失败
|
||||
> - build_env.sh 超时 720 秒或返回码非 0 → 评测失败
|
||||
|
||||
## 附录 B：相关资源
|
||||
|
||||
| 资源 | 链接 |
|
||||
|------|------|
|
||||
| GRAB 论文 | arXiv 2602.01865 |
|
||||
| HSTU 论文 | arXiv 2402.17152 (ICML 2024) |
|
||||
| 官方 Baseline 项目 | AI Studio project/10186630 |
|
||||
| 比赛主页 | aistudio.baidu.com/competition/detail/1461 |
|
||||
| 提交结果页 | aistudio.baidu.com/competition/detail/1461/0/submit-result |
|
||||
Reference in New Issue
Block a user