Transformers Trainer怎样保留自定义字段传入模型

如果你在 Hugging Face Transformers 的 Trainer 中加入了 sample_weight、group_id 之类的自定义列，却发现它们没有出现在模型输入里，先检查 remove_unused_columns。Trainer 默认会根据模型 forward() 的参数名裁剪数据集列；要保留字段，通常有两条路：关闭这一步裁剪，并让模型能接收该字段；或者直接把字段写进自定义模型的 forward() 签名。

官方地址：https://huggingface.co/docs/transformers/main_classes/trainer

排查的关键不是“数据集有没有这一列”，而是“Trainer 准备 batch 时认为模型需要哪些参数”。字段在数据集里存在，不等于它会自动传给模型。

用Transformers库的Trainer做模型微调时，默认只会把模型forward方法入参声明过的字段传入，想要自己新增的自定义字段也能顺利传到模型内部，只要调整数据处理、训练流程的对应环节适配就能实现。

不需要修改Trainer核心源码，优先通过自定义数据整理函数过滤保留需要的自定义字段，再配合覆写compute_loss或者训练步骤的逻辑，就能把自定义字段透传给模型。

先确认字段在哪一步消失

先把问题拆成三层：数据集原始列、Trainer 组装后的 batch，以及模型 forward() 实际接收的关键字参数。官方文档说明，传入普通 Dataset 时，不被 model.forward() 接受的列会被自动移除。假设数据集有 input_ids、attention_mask、labels 和 sample_weight，而模型签名只有前三个参数，那么 sample_weight 可能在进入 data collator 或模型之前就不见了。

可以先用一个最小的 batch 检查列名，但不要把“打印数据集列名”当成最终结论。真正有用的对照是：模型参数列表里有没有同名参数，以及训练参数是否允许未使用列继续流转。

改法一：关闭自动列裁剪

如果你希望自定义字段完整保留到 data collator 和模型调用阶段，可以在 TrainingArguments 中设置 remove_unused_columns=False。但这个选项只负责“不删除”，不会替你修改模型签名。模型仍然要能接收传入的关键字，否则下一步会变成 unexpected keyword argument。

from transformers import TrainingArguments

# 保留数据集中的自定义列，后续由模型决定如何处理
training_args = TrainingArguments(
    output_dir="./runs/weighted",
    remove_unused_columns=False,
    per_device_train_batch_size=8,
)

# 数据集中的 sample_weight 会继续留在 batch 中
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)

这条路适合 data collator 还要使用额外列，或者模型接口本来就准备接收多个扩展字段的场景。代价是所有列都会更容易流到模型调用处，字段命名、类型和清理责任要由你自己承担。

改法二：让 forward 显式接收自定义字段

如果只想让某个字段进入模型，通常更清晰的方式是把它加入自定义模型的 forward()。Trainer 会根据参数名识别它，列裁剪仍然可以保持默认行为。下面用加权分类损失示意，sample_weight 被声明为可选参数，然后在模型内部参与 loss 计算。

import torch
import torch.nn.functional as F
from transformers.modeling_outputs import SequenceClassifierOutput

class WeightedClassifier(torch.nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model

    def forward(self, input_ids=None, attention_mask=None,
                labels=None, sample_weight=None):
        # 只把基础模型需要的标准字段传下去，避免额外字段继续扩散
        outputs = self.base_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        if labels is None:
            # 预测阶段没有标签时，直接返回 logits
            return SequenceClassifierOutput(logits=outputs.logits)

        # 逐样本计算损失，便于应用自定义权重
        per_item_loss = F.cross_entropy(outputs.logits, labels, reduction="none")
        if sample_weight is None:
            loss = per_item_loss.mean()
        else:
            # 权重要放到同一设备，并防止权重总和为零
            weights = sample_weight.to(per_item_loss.device).float()
            loss = (per_item_loss * weights).sum() / weights.sum().clamp_min(1e-8)
        return SequenceClassifierOutput(loss=loss, logits=outputs.logits)

这里的重点是参数名完全匹配：数据集列叫 sample_weight，forward() 也叫 sample_weight。若列名是 weight，就要同步调整签名或在数据整理阶段完成一次明确的重命名，不要依赖位置顺序猜测。

按现象排查配置和签名

字段消失和字段报错是两个不同阶段的问题。字段消失通常说明 Trainer 仍在按默认规则裁剪；字段报错则常见于关闭裁剪后，模型没有对应参数。可以按下面的顺序排查：

实际选择上，如果扩展字段只服务一个明确模型，优先显式写进 forward()；如果字段还要经过自定义 collator 或多个模型层级，才考虑关闭自动列裁剪。无论采用哪条路，都要让数据列名、collator 输出键和 forward 参数保持一致。

常见问题

只设置 remove_unused_columns=False 就一定能传入字段吗？

不一定。它只是不再提前删除列；如果模型不接受这个关键字，调用时仍会报参数错误。还需要扩展模型签名，或者在 collator 中把该字段转换为模型已经支持的输入。

可以把 sample_weight 写成 **kwargs 接收吗？

可以作为兜底，但不如显式参数直观。Trainer 是按签名判断列是否被使用的，明确写出 sample_weight=None 更容易让列裁剪保留它，也更方便读代码的人确认字段用途。

参考地址：https://huggingface.co/docs/transformers/main_classes/trainer