第三名解决方案

祝贺所有获奖者！感谢组织者举办如此有趣的比赛。LMSYS 是我第一次获得银牌的 LLM 竞赛。这次我非常高兴能在 WSDM Cup 获得金牌。在这里让我分享我的解决方案。

简要总结 (TL;DR)

• 重新实现了 Eedi Rerank Model 的训练 pipeline。使用 AutoModelForCasualLM 而不是 AutoModelForSequenceClassification 以便通过 vllm 进行更快的推理。
• Qwen/Qwen2.5-14B-Instruct (带有后预训练) + Phi4 (无后预训练)。合并了使用不同种子训练的模型。
• 像 LMSYS 第一名 那样使用 72B 模型进行蒸馏。但我发现 72B 模型并不是必须的。使用 14B 模型的 logit 并执行“自蒸馏”可以达到相同的 CV 分数。也许是软 logit 带来的标签清洗效果真正起了作用。
• 使用 auto-round 进行量化。
• vllm 推理。对于集成，首先按 token 长度排序，然后对 25% 的数据使用 Qwen2.5-14B 进行 TTA。剩余的推理时间用于 Phi4。

预处理

使用 fasttext 推断 prompt 的语言。

训练部分

我重写了 SFTTrainer 的 compute_loss 函数，仅计算 Token "A" 和 "B" 的 loss。

class SFTChoiceTrainer(SFTTrainer):
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None, return_choice_logit=False):
        labels = inputs['labels']
        labels[labels > 57] = -100
        inputs['labels'] = labels
        _, outputs = super().compute_loss(model, inputs, return_outputs=True, num_items_in_batch=num_items_in_batch)
        logits = outputs.logits
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        logits_target = []
        labels_target = []
        for i in range(len(shift_labels)):
            lbl = shift_labels[i].cpu().numpy()
            target_idx = np.where(lbl != -100)[0][0]
            # token id for "A" and "B" is 32 and 33
            logits_target.append(shift_logits[i][target_idx][32:34])
            labels_target.append(shift_labels[i][target_idx]-32)

        # (batch_size, 2)
        logits_target = torch.stack(logits_target, dim=0)
        # (batch_size)
        labels_target = torch.tensor(labels_target).to(outputs.logits.device)
        loss = F.cross_entropy(logits_target, labels_target)
        if return_choice_logit:
            return (loss, outputs, logits_target) if return_outputs else (loss, logits_target)
        else:
            return (loss, outputs) if return_outputs else loss

prompt、response a 和 response b 的最大长度限制为 2048。因此输入的最大长度约为 6144。

class PromptManager(object):
    preference_prompt_predict_template = '''<| ( performance|>system
You are a highly skilled AI assistant. You will be provided a request from user and several responses from different assistants. Your job is to judge which response is the best. Only answer the letter of the best response.<| ( eyes|>
<| ( performance|>user
<Request>
{prompt}
</Request>

<Language>
{language}
</Language>

{responses}<| ( eyes|>
<| ( performance|>assistant
'''

    preference_prompt_train_template = preference_prompt_predict_template + '{answer}<| ( eyes|>\n'

    response_template = '''<Response_{choice}>
{response}
</Response_{choice}>\n'''

    sep = '<| ( performance|>assistant\n'

后预训练 (Post-Pretrain)

我使用了以下数据进行后预训练。后预训练与训练共享相同的训练 pipeline。

• ultrafeedback
• LMSYS 第二名 收集的数据集

训练数据

我使用了以下数据进行训练。

• LMSYS 数据
• LMSYS 33k
• WSDM Cup 数据

对于教师模型 (Qwen2.5-72B 和 Athene-V2-Chat)，我在新年假期期间使用 deepspeed 在多 GPU 上进行训练。为了避免调整学习率，我只需将 gradient_accumulation_steps 设置为 1，以保持与单 GPU 训练相同的 batch size。

交叉验证 (无蒸馏)

基本上，14B 模型的学习率 10e-5 是好的。Loss spike 开始在 14e-5 左右出现。对于 72B 模型，较小的学习率更好。

交叉验证 (含蒸馏)

对我来说最好的 Temperature 是 5.0 (实验 2~5)，最好的 soft loss 权重是 0.9 (实验 2~7)。
遵循 LMSYS 第一名的做法，我尝试平均 KLDivLoss 和 CosineEmbeddingLoss，但效果不明显。KLDivLoss 已经足够。

量化 (Quantize)

我使用 auto-round 进行量化。使用较重的量化设置可以抑制准确率下降。量化实验基于 Qwen2.5 3B 模型。
我最终的量化设置是 seqlen=6144, nsamples=512, iters=3000。

推理部分

我使用 vllm 进行推理。推理部分没有什么特别的，全是工程优化。我使用了以下技巧。

• 根据测试数据大小动态改变推理时间限制。对于训练阶段，推理时间限制设置为 265 分钟。对于预测阶段，推理时间限制设置为 690 分钟。
• 使用 T4x2 环境测试模型的吞吐量。模型的平均吞吐量约为每分钟 70000 token。对于 TTA，得益于 vllm 的前缀缓存，吞吐量可以达到每分钟约 74000 token。
• 按 token 长度对测试数据进行排序，并对前 25% 的数据使用 Qwen2.5 14B 模型进行 TTA。
• 监控剩余的推理时间，并根据吞吐量计算 Phi4 可以处理的 token 数量。在训练阶段大约可以处理 60% 的测试数据。

提交策略

基本上，我选择排行榜分数最高的提交。由于排行榜分数不稳定，我使用不同的种子训练了相同的模型并线性合并了模型权重。
我发现 Phi4 对短 prompt 表现良好，但在排行榜上对长 prompt 表现不佳。所以我根据 token 长度改变了 Phi4 的集成权重。

1. 计算测试数据 prompt 长度的分位数。q1=0.25, q2=0.5, q3=0.75, q4=1.0
2. 对于 token 长度 <= q2，Phi4 权重设置为 0.5。
3. 对于 token 长度 > q2 且 <= q3，Phi4 权重设置为 0.15。
4. 对于 token 长度 > q3，Phi4 权重设置为 0.1。(由于推理时间限制未使用。Phi4 在训练阶段只能推理约 60% 的测试数据。)

尝试过但放弃的方法

• 受 Eedi 第一名启发，我蒸馏了 Qwen2.5-3b 来生成关于如何解决 prompt (=用户请求) 的 CoT。CV 和 LB 提高了 0.001~0.002，但推理时间太长，我经常遇到超时错误。
• 尝试了由 deepseek 蒸馏的 Qwen2.5-14b 和 Qwen2.5-14b-1m。LB 效果不好。

# CoT
system_prompt_template = 'You are an excellant AI assistant. You will be given a prompt from the user. Your task is to infer the task user request you to do.'

user_prompt_template = '''### Prompt Start ###
{text}
### Prompt End ###
Please answer the task user request you to do with one sentence in English and then explain the key factors to complete the task step by step briefly. No need to provide the answer for the prompt.
'''