CTI-Inference-Opt/推理优化方案.md

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 占位：

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 签名

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 函数中的模型加载部分：

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 中同步处理输入数据精度：

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 转换：

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 / 标准实现）：

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：

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() 之后添加一行：

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 中用一个小输入做一次预热编译：

#!/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() 的数据加载部分，缓存时就完成设备搬运和类型转换：

# 原代码
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：

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 当前状态

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 的被激活次数：

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：

# 仅当某层 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），有精度损失风险。

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 设计

#!/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