Trainer自定义训练数据的顺序

transformer库的Trainer自定义训练数据的顺序

背景为在使用 transformer 库来 SFT 一个大模型时，需要控制训练数据的有序输入，而在Trainer 中并没有相应的参数实现这一需求。

说明：

• 此处的训练采用了torch.distributed.launch 来进行训练，即启动训练的命令应该是下面这个样子的: CUDA_VISIBLE_DEVICES=0,1,2 python -m torch.distributed.launch --master_addr ${MASTER_ADDR} --master_port ${MASTER_PORT} --nproc_per_node=${nproc_per_node} --use_env my_train_math.py。
• 通过添加参数 --shuffle False 的方法阻止dataloader 打乱数据的排序，然而实际使用时发现训练数据依旧是无序输入的。

解决方法

定位原因

负责控制取 batch 数据的是Data loader，因此查看Trainer 的源代码，可以看到get_train_dataloader() 这个函数。在Dataloader的构建中，负责数据分发任务的是sampler，直接看第28行，查找self._get_train_sampler() 的具体实现。

def get_train_dataloader(self) -> DataLoader:
    """
        Returns the training [`~torch.utils.data.DataLoader`].

        Will use no sampler if `train_dataset` does not implement `__len__`, a random sampler (adapted to distributed
        training if necessary) otherwise.

        Subclass and override this method if you want to inject some custom behavior.
        """
    if self.train_dataset is None:
        raise ValueError("Trainer: training requires a train_dataset.")

        train_dataset = self.train_dataset
        data_collator = self.data_collator
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="training")

            dataloader_params = {
                "batch_size": self._train_batch_size,
                "collate_fn": data_collator,
                "num_workers": self.args.dataloader_num_workers,
                "pin_memory": self.args.dataloader_pin_memory,
            }

            if not isinstance(train_dataset, torch.utils.data.IterableDataset):
                dataloader_params["sampler"] = self._get_train_sampler()	# 只有这一行需要关注
                dataloader_params["drop_last"] = self.args.dataloader_drop_last
                dataloader_params["worker_init_fn"] = seed_worker

                return self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))

如下是_get_train_sampler() 的具体实现，可以看到里面其实最后能够用到的sampler 其实只可能是RandomSampler 和 LengthGroupedSampler，如果自定义输入顺序的需求无法通过LengthGroupedSampler 满足，就得换用别的sampler进行替换。

此时找到了原因，即Trainer 内部只提供RandomSampler 和 LengthGroupedSampler。

def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]:
    if self.train_dataset is None or not has_length(self.train_dataset):
        return None

    # Build the sampler.
    if self.args.group_by_length:
        if is_datasets_available() and isinstance(self.train_dataset, datasets.Dataset):
            lengths = (
                self.train_dataset[self.args.length_column_name]
                if self.args.length_column_name in self.train_dataset.column_names
                else None
            )
        else:
            lengths = None
            model_input_name = self.tokenizer.model_input_names[0] if self.tokenizer is not None else None
            return LengthGroupedSampler(
                self.args.train_batch_size * self.args.gradient_accumulation_steps,
                dataset=self.train_dataset,
                lengths=lengths,
                model_input_name=model_input_name,
            )

        else:
            return RandomSampler(self.train_dataset)

解决方案

为了自定义输入顺序，有如下两个解决方案：

1. 自己构建sampler，在sampler的item() 中自行实现有序输入。

2. 调整数据集的顺序，然后使用顺序输出的sampler（SequentialSampler），组合之下实现自定义输入顺序。

我的方法比较讨巧。因为怕在训练过程中引入奇奇怪怪的bug，没有选择自行构建sampler，因此用了方法2。

• 首先是调整数据集的顺序，直接魔改 getitem__(self, i) 就好。

• 顺序输出训练数据可以采用torch.utils.data.SequentialSampler 实现，其会从前到后有序的进行数据的分发，文档可见：torch.utils.data.sampler — PyTorch 2.4 documentation。

个人感觉最简单的方法就是自己写一个Train 的子类，然后重 写_get_train_sampler() 函数，返回自己需要的函数。最后反映到代码中的改动很小：

# 构建自定义的Trainer类，重写_get_train_sampler，在不shuffle的情况下无脑返回SequentialSampler
class CustomTrainer(Trainer):
    def _get_train_sampler(self):
        if self.train_dataset is None or not hasattr(self.train_dataset, '__len__') or len(self.train_dataset) == 0:
            return None

        # 如果你想根据某种条件来选择不同的 Sampler，比如按照顺序加载或者随机采样
        if self.args.shuffle:
            return RandomSampler(self.train_dataset)
        else:
            # 使用 SequentialSampler 保证数据顺序一致
            print("Using SequentialSampler !!!!!!!!!!!!!!!")
            return SequentialSampler(self.train_dataset)
        
# 使用CustomTrainer 的方法与Trainer一致：
trainer = CustomTrainer(
    model=model, 
    tokenizer=tokenizer, 
    args=training_args, 
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    data_collator=data_collator
)
trainer.train()