Prhub

#2237 fix: keep dataset order in filter_long_prompt for mixed multimodal data

原始 PR 作者 keepkeen 合并时间 2026-08-12 13:35 文件变更 4 提交数 2 评论 0 代码增减 +132 / -9

执行摘要

修复 filter_long_prompt 混合多模态数据顺序错乱

PR body 指出:'With a processor configured, filter_long_prompt reorders the dataset — even when it filters nothing out.' 示例输入 p0..p5 输出为 p1,p3,p5,p0,p2,p4。由于 --rollout-shuffle 默认为 False,该顺序即训练顺序,导致 'A VLM run over a mixed text/image dataset trains every text-only prompt before it ever sees an image.' 这是 #1662 吞吐优化引入的回归,本 PR 旨在保留优化同时恢复正确顺序。

值得精读:修复点集中、测试设计精巧(用替身与桩避免引入 transformers),是优化引入回归后最小代价恢复正确性的典型样例。数据管线与 VLM 训练相关同学建议细看 filter_long_prompt 的实现与测试写法。

讨论亮点

该 PR 无公开 review 评论。PR body 中作者说明了设计权衡:分组评分是 #1662 引入的吞吐优化,不能因为修复顺序而放弃;因此采用携带原始位置 + 结束后排序的方式,既保留批量 tokenizer 的高吞吐,又让返回顺序与原始数据集一致。代码注释也强调 'The two groups are scored separately for throughput, so restore the dataset order here'。

实现拆解

  1. 在 slime/utils/data.py 的 filter_long_prompt 中,将分组循环改为 for position, sample in enumerate(origin_samples),text_only 与 multimodal 列表元素变为 (position, sample) 元组,为后续恢复顺序保留原始位置信息。
  2. 将原 filtered_samples 改名为 kept,两组分支中追加 (position, sample);两组处理完后用 kept.sort(key=lambda position_and_sample: position_and_sample[0]) 按原始位置排序,再解包成 filtered_samples 返回。排序只在数据集构造时执行一次,不在热路径上。
  3. 无 processor 的 else 分支天然按 origin_samples 顺序处理,无需修改。
  4. 新增 tests/test_filter_long_prompt.py,使用 _Tokenizer/_Processor 替身和 stub_process_vision_info 桩模块,覆盖无过滤保持顺序、混合过滤保持顺序、单模态批量和无 processor 路径四类场景,其中两类在 main 上会失败。
  5. 在 .github/workflows/pr-test.yml.j2 模板及生成的 pr-test.yml 中将 test_filter_long_prompt.py 注册到 cpu-unittest 作业(num_gpus: 0),确保 CI 覆盖。
文件 模块 状态 重要度
slime/utils/data.py 数据过滤 modified 6.32
tests/test_filter_long_prompt.py 数据过滤 added 7.2
.github/workflows/pr-test.yml 持续集成 modified 2.73
.github/workflows/pr-test.yml.j2 持续集成 modified 2.24

关键符号

filter_long_prompt stub_process_vision_info test_preserves_order_when_nothing_is_filtered test_preserves_order_when_some_are_filtered test_single_modality_batches_are_unchanged test_no_processor_path_still_preserves_order

关键源码片段

slime/utils/data.py core-logic

核心修复文件:filter_long_prompt 在 processor 分支中通过携带原始位置并排序,恢复混合多模态数据集的原始顺序。

def filter_long_prompt(origin_samples: list[Sample], tokenizer, processor, max_length: int | None) -> list[Sample]:
    # max_length 为 None 时直接返回,由调用方决定是否过滤。
    if max_length is None:
        return origin_samples
​
    # prompt 不是字符串(如 list 形式)时跳过过滤,并提示配置 apply_chat_template。
    if not isinstance(origin_samples[0].prompt, str):
        logger.warning(
            'Skipping max_length check for list prompt. '
            'Set apply_chat_template=True to enable length filtering.'
        )
        return origin_samples
​
    if processor:
        # 为了吞吐量,把样本分为纯文本和多模态两组:
        # 纯文本用批量 tokenizer 一次打分,多模态用 processor 逐条打分。
        # 分组时同时记录原始 position,便于之后恢复数据集顺序。
        text_only = []
        multimodal = []
        for position, sample in enumerate(origin_samples):
            if sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()):
                multimodal.append((position, sample))
            else:
                text_only.append((position, sample))
​
        kept = []
        if text_only:
            # 纯文本组:一次调用 tokenizer,得到所有样本的 input_ids。
            prompts = [s.prompt for _, s in text_only]
            input_ids_list = tokenizer(prompts, add_special_tokens=False)['input_ids']
            for (position, sample), input_ids in zip(text_only, input_ids_list, strict=True):
                if len(input_ids) <= max_length:
                    kept.append((position, sample))
​
        if multimodal:
            # 多模态组:processor 需要按样本提取视觉信息,只能逐个处理。
            from slime.utils.processing_utils import process_vision_info
            for position, sample in multimodal:
                multimodal_inputs = process_vision_info(sample.prompt, processor)
                processor_output = processor(text=sample.prompt, **multimodal_inputs)
                input_ids = processor_output['input_ids'][0]
                if len(input_ids) <= max_length:
                    kept.append((position, sample))
​
        # 两组分开评分只是优化,必须按原始位置排序恢复数据集顺序;
        # 否则未开启 --rollout-shuffle 时,训练会先遍历纯文本,再遍历多模态。
        kept.sort(key=lambda position_and_sample: position_and_sample[0])
        filtered_samples = [sample for _, sample in kept]
    else:
        # 没有 processor 时所有样本都走批量 tokenizer,天然保持原始顺序。
        prompts = [sample.prompt for sample in origin_samples]
        input_ids_list = tokenizer(prompts, add_special_tokens=False)['input_ids']
        filtered_samples = [
            sample
            for sample, input_ids in zip(origin_samples, input_ids_list, strict=True)
            if len(input_ids) <= max_length
        ]
​
    logger.info(
        f'Filtered {len(origin_samples) - len(filtered_samples)} samples longer than max_length={max_length}.'
    )
    return filtered_samples
tests/test_filter_long_prompt.py test-coverage

新增 CPU 单元测试,覆盖顺序保持的核心场景,锁定回归行为。

# slime.utils.processing_utils 会拉起 transformers 和 PIL,
# CPU 测试镜像没有这些依赖;多模态分支只用到 process_vision_info,
# 所以这里用 types.ModuleType 构造一个假模块桩掉它。
@pytest.fixture
def stub_process_vision_info(monkeypatch):
    import sys
    import types
    module = types.ModuleType('slime.utils.processing_utils')
    module.process_vision_info = lambda prompt, processor: {'images': None}
    monkeypatch.setitem(sys.modules, 'slime.utils.processing_utils', module)
​
​
class _Tokenizer:
    # 批处理 tokenizer 替身:prompt 'pN:len' 编码为 len 个 id。
    def __call__(self, prompts, add_special_tokens=False):
        return {'input_ids': [[0] * _encoded_length(p) for p in prompts]}
​
​
class _Processor:
    # 逐样本 processor 替身,长度约定与 _Tokenizer 一致。
    def __call__(self, text=None, **kwargs):
        return {'input_ids': [[0] * _encoded_length(text)]}
​
​
def _encoded_length(prompt: str) -> int:
    # prompt 形如 'pN:len',冒号后即编码长度,用整型解析出来。
    return int(prompt.split(':')[1])
​
​
@pytest.mark.unit
def test_preserves_order_when_some_are_filtered(stub_process_vision_info):
    # 混合纯文本与多模态样本,部分超长被过滤。
    # 规格为 ( 是否多模态 , 编码长度 ),max_length=100。
    specs = [
        (True, 5), # p0 多模态,保留
        (False, 500), # p1 纯文本,丢弃
        (False, 5), # p2 纯文本,保留
        (True, 500), # p3 多模态,丢弃
        (False, 5), # p4 纯文本,保留
        (True, 5), # p5 多模态,保留
    ]
    samples = _make_samples(specs) # _make_samples 按 specs 构造 Sample 列表
​
    kept = filter_long_prompt(samples, _Tokenizer(), _Processor(), max_length=100)
​
    # 关键断言:幸存样本必须保持原始数据集顺序,而不是先纯文本后多模态。
    assert [s.prompt for s in kept] == ['p0:5', 'p2:5', 'p4:5', 'p5:5']

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

  1. 行为变更:修复会让依赖旧顺序(纯文本在前)的实验结果变化,这是修复目标,但对比实验需注意基线差异。
  2. 边界条件:filter_long_prompt 对空列表仍会因 origin_samples[0] 抛 IndexError,本次未处理,延续原有行为。
  3. 性能:排序 O(n log n) 仅在数据集构造时执行一次,可忽略;分组与 scoring 逻辑不变。
  4. 稳定性:kept.sort 使用稳定排序且 position 唯一,结果确定;无 processor 分支不受影响。

影响面:所有配置了 processor 且数据混合文本与多模态的训练 run,训练顺序从纯文本全在前修正为原始数据顺序,更符合数据分布,利于 VLM 训练稳定。对纯文本数据集或未配置 processor 的场景零影响。接口、配置项、schema 无变化;CI 增加一个 CPU 测试,团队维护成本极小。

训练顺序行为变更 空数据集边界未处理

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论