通义千问3-Reranker-0.6B训练指南：自定义数据集fine-tuning

1. 引言

如果你正在构建智能搜索、推荐系统或者RAG应用，可能会遇到这样的问题：检索出来的结果很多，但真正相关的却没几个。通义千问3-Reranker-0.6B就是专门解决这个痛点的模型，它能够对初步检索结果进行精细排序，把最相关的内容排到最前面。

但现成的模型可能不完全适合你的特定场景和数据类型。这就是为什么需要自定义训练的原因——通过fine-tuning，你可以让模型更好地理解你的领域术语、业务逻辑和用户需求。

本文将手把手带你完成整个训练流程，从数据准备到模型评估，即使你是刚接触reranker的新手，也能跟着做下来。

2. 环境准备与快速部署

2.1 系统要求与依赖安装

首先确保你的环境满足基本要求：

• Python 3.8或更高版本
• PyTorch 2.0+
• CUDA 11.7或更高版本（如果使用GPU）
• 至少16GB内存（推荐32GB）

安装必要的依赖包：

pip install transformers>=4.51.0
pip install datasets
pip install accelerate
pip install peft
pip install sentencepiece
pip install protobuf

2.2 模型下载与初始化

从Hugging Face下载Qwen3-Reranker-0.6B模型：

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "Qwen/Qwen3-Reranker-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
model = AutoModelForCausalLM.from_pretrained(model_name)

# 检查模型参数
print(f"模型参数量: {model.num_parameters():,}")

3. 数据准备与格式化

3.1 理解reranker训练数据格式

Reranker训练需要三元组数据：(查询, 正例文档, 负例文档)。正例文档与查询高度相关，负例文档相关性较低或无关。

# 示例数据格式
training_examples = [
    {
        "query": "如何安装Python环境",
        "positive": "Python安装教程：从官网下载安装包，运行安装程序，勾选添加PATH选项",
        "negative": "Python是一种高级编程语言，由Guido van Rossum创建"
    },
    # 更多数据...
]

3.2 构建自定义数据集

根据你的业务场景准备数据。如果你有用户点击日志、人工标注数据或相关性评分，可以这样处理：

from datasets import Dataset

def prepare_reranker_data(raw_data):
    processed_data = []
    for item in raw_data:
        # 假设raw_data包含query, positive_doc, negative_docs列表
        for negative_doc in item["negative_docs"]:
            processed_data.append({
                "query": item["query"],
                "positive": item["positive_doc"],
                "negative": negative_doc
            })
    return Dataset.from_list(processed_data)

# 加载你的数据
your_data = load_your_custom_data()  # 替换为你的数据加载逻辑
dataset = prepare_reranker_data(your_data)

3.3 数据预处理与tokenization

Reranker模型有特定的输入格式要求，需要正确格式化：

def format_reranker_input(query, document, instruction=None):
    if instruction is None:
        instruction = "给定查询和文档，判断文档是否与查询相关"
    
    return f"<|im_start|>system\n基于查询和指令判断文档相关性，只能回答'yes'或'no'。<|im_end|>\n<|im_start|>user\n<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {document}<|im_end|>\n<|im_start|>assistant\n"

def tokenize_function(examples):
    batch_size = len(examples["query"])
    processed_texts = []
    
    for i in range(batch_size):
        # 正例样本
        positive_text = format_reranker_input(
            examples["query"][i], 
            examples["positive"][i]
        )
        # 负例样本
        negative_text = format_reranker_input(
            examples["query"][i], 
            examples["negative"][i]
        )
        processed_texts.extend([positive_text, negative_text])
    
    # Tokenize
    tokenized = tokenizer(
        processed_texts,
        truncation=True,
        padding=False,
        max_length=8192,
        return_tensors=None
    )
    
    return tokenized

# 应用tokenization
tokenized_dataset = dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=dataset.column_names
)

4. 训练参数设置与优化

4.1 基础训练配置

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./qwen3-reranker-finetuned",
    learning_rate=2e-5,
    per_device_train_batch_size=2,  # 根据GPU内存调整
    per_device_eval_batch_size=2,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir="./logs",
    logging_steps=10,
    save_steps=500,
    eval_steps=500,
    evaluation_strategy="steps",
    save_total_limit=2,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    fp16=True,  # 启用混合精度训练
)

4.2 自定义损失函数

由于reranker训练是对比学习，需要自定义损失函数：

import torch
import torch.nn as nn
import torch.nn.functional as F

class RerankerLoss(nn.Module):
    def __init__(self, token_true_id, token_false_id):
        super().__init__()
        self.token_true_id = token_true_id
        self.token_false_id = token_false_id
    
    def forward(self, logits, labels):
        # 获取每个序列最后一个token的logits
        last_token_logits = logits[:, -1, :]
        
        # 提取"yes"和"no"的logits
        true_logits = last_token_logits[:, self.token_true_id]
        false_logits = last_token_logits[:, self.token_false_id]
        
        # 计算二元分类概率
        batch_scores = torch.stack([false_logits, true_logits], dim=1)
        probabilities = F.log_softmax(batch_scores, dim=1)
        
        # 计算对比损失
        positive_scores = probabilities[::2, 1]  # 正例的"yes"概率
        negative_scores = probabilities[1::2, 1]  # 负例的"yes"概率
        
        # 希望正例得分高于负例得分
        loss = -torch.log(torch.sigmoid(positive_scores - negative_scores)).mean()
        return loss

# 获取特殊token的ID
token_true_id = tokenizer.convert_tokens_to_ids("yes")
token_false_id = tokenizer.convert_tokens_to_ids("no")
loss_fn = RerankerLoss(token_true_id, token_false_id)

4.3 创建Trainer实例

from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.get("labels")
        outputs = model(**inputs)
        logits = outputs.logits
        
        loss = loss_fn(logits, labels)
        return (loss, outputs) if return_outputs else loss

trainer = CustomTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    eval_dataset=tokenized_dataset,  # 实际使用时应该划分验证集
    tokenizer=tokenizer,
)

5. 开始训练与监控

5.1 启动训练过程

# 开始训练
train_result = trainer.train()

# 保存最终模型
trainer.save_model()
tokenizer.save_pretrained(training_args.output_dir)

# 保存训练指标
metrics = train_result.metrics
trainer.log_metrics("train", metrics)
trainer.save_metrics("train", metrics)

5.2 训练过程监控

训练过程中关注这些指标：

• 训练损失：应该持续下降
• 验证损失：避免过拟合，应该与训练损失同步下降
• GPU内存使用：确保没有内存溢出
• 训练速度：监控每秒处理的样本数

如果发现过拟合，可以：

• 增加权重衰减(weight_decay)
• 减少训练轮数(epochs)
• 增加dropout比例
• 使用早停(early stopping)

6. 模型评估与测试

6.1 创建评估函数

def evaluate_reranker(model, tokenizer, test_data):
    model.eval()
    correct = 0
    total = 0
    
    for example in test_data:
        # 准备正例输入
        positive_input = format_reranker_input(example["query"], example["positive"])
        positive_tokens = tokenizer(positive_input, return_tensors="pt")
        
        # 准备负例输入
        negative_input = format_reranker_input(example["query"], example["negative"])
        negative_tokens = tokenizer(negative_input, return_tensors="pt")
        
        with torch.no_grad():
            # 计算正例得分
            positive_output = model(**positive_tokens)
            positive_logits = positive_output.logits[:, -1, :]
            positive_score = torch.softmax(positive_logits[:, [token_false_id, token_true_id]], dim=1)[0, 1]
            
            # 计算负例得分
            negative_output = model(**negative_tokens)
            negative_logits = negative_output.logits[:, -1, :]
            negative_score = torch.softmax(negative_logits[:, [token_false_id, token_true_id]], dim=1)[0, 1]
        
        # 检查正例得分是否高于负例得分
        if positive_score > negative_score:
            correct += 1
        total += 1
    
    accuracy = correct / total
    return accuracy

# 加载测试数据
test_dataset = load_test_data()  # 替换为你的测试数据加载逻辑
accuracy = evaluate_reranker(model, tokenizer, test_dataset)
print(f"模型准确率: {accuracy:.2%}")

6.2 实际应用测试

测试模型在你的实际场景中的表现：

def rerank_documents(query, documents, model, tokenizer, top_k=3):
    """
    对文档列表进行重排序
    """
    scores = []
    
    for doc in documents:
        input_text = format_reranker_input(query, doc)
        inputs = tokenizer(input_text, return_tensors="pt")
        
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits[:, -1, :]
            doc_score = torch.softmax(logits[:, [token_false_id, token_true_id]], dim=1)[0, 1].item()
        
        scores.append((doc, doc_score))
    
    # 按得分降序排序
    scores.sort(key=lambda x: x[1], reverse=True)
    return scores[:top_k]

# 示例使用
query = "如何学习深度学习"
documents = [
    "深度学习是机器学习的一个分支，使用神经网络",
    "深度学习入门教程：从基础概念到实践项目",
    "Python编程基础语法介绍",
    "深度学习框架比较：TensorFlow vs PyTorch"
]

ranked_results = rerank_documents(query, documents, model, tokenizer)
print("重排序结果:")
for i, (doc, score) in enumerate(ranked_results):
    print(f"{i+1}. 得分: {score:.3f} - {doc[:50]}...")

7. 常见问题与解决方案

7.1 内存不足问题

如果遇到GPU内存不足：

# 解决方案1：减小batch size
training_args.per_device_train_batch_size = 1

# 解决方案2：使用梯度累积
training_args.gradient_accumulation_steps = 4

# 解决方案3：使用LoRA等参数高效微调方法
from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.1,
    bias="none",
    task_type="CAUSAL_LM"
)
model = get_peft_model(model, lora_config)

7.2 过拟合处理

如果模型在训练集上表现很好但在测试集上差：

• 增加训练数据量
• 使用数据增强技术
• 添加正则化（dropout、权重衰减）
• 早停策略

7.3 训练不稳定

如果损失波动很大：

• 减小学习率
• 使用学习率预热
• 添加梯度裁剪

training_args.learning_rate = 1e-5
training_args.warmup_steps = 100
training_args.max_grad_norm = 1.0

8. 总结

通过这篇教程，你应该已经掌握了如何对通义千问3-Reranker-0.6B进行自定义数据集的fine-tuning。整个过程从环境准备开始，涵盖了数据格式化、训练参数配置、模型训练到最终评估的完整流程。

实际使用中，最关键的是准备好高质量的训练数据——好的数据能让模型学会真正理解你的业务场景和相关性标准。训练过程中要多关注验证集的表现，避免过拟合。如果资源有限，可以考虑使用LoRA等参数高效微调方法。

训练完成后，别忘了在实际场景中全面测试模型效果，看看重排序后的结果是否真的更符合用户需求。有时候可能需要多次迭代调整训练数据和参数，才能达到最佳效果。

---

获取更多AI镜像

想探索更多AI镜像和应用场景？访问 CSDN星图镜像广场，提供丰富的预置镜像，覆盖大模型推理、图像生成、视频生成、模型微调等多个领域，支持一键部署。