fix: keep dataset order in filter_long_prompt for mixed multimodal data

[keepkeen]

What

With a processor configured, filter_long_prompt reorders the dataset — even when it filters nothing out.

Root cause

The function scores text-only samples with one batched tokenizer call and multimodal samples one at a time through the processor, then appends the survivors group by group:

filtered_samples = []
if text_only:
    ...
    filtered_samples.append(sample)     # all text-only survivors first
if multimodal:
    ...
    filtered_samples.append(sample)     # then all multimodal survivors

So the returned list is every surviving text-only sample followed by every surviving multimodal sample, regardless of where they sat in the dataset.

Dataset.__init__ assigns that list to self.origin_samples / self.samples, and --rollout-shuffle defaults to False, so this is the training order. A VLM run over a mixed text/image dataset trains every text-only prompt before it ever sees an image.

Reproduction

Six alternating samples, max_length high enough that nothing is dropped:

input order : ['p0', 'p1', 'p2', 'p3', 'p4', 'p5']
output order: ['p1', 'p3', 'p5', 'p0', 'p2', 'p4']
nothing filtered? True

Why it's a regression

Before #1662 this was a single order-preserving loop over origin_samples. That PR split it into two groups for throughput ("Use processor only for samples with actual multimodal content; use batched tokenizer for text-only") and lost the ordering as a side effect — the stated goal was speed, not a different sample order.

Fix

Keep the split, carry each sample's original position through it, and sort the survivors back before returning. filter_long_prompt runs once at dataset construction, so the sort is not on any hot path.

Test

New CPU unit test tests/test_filter_long_prompt.py, registered in the cpu-unittest job: order is preserved with nothing filtered, with a mix of dropped text-only and dropped multimodal samples, for single-modality batches, and on the no-processor path. Two of the five cases fail on main.

slime.utils.processing_utils is stubbed in the test — the multimodal branch imports it lazily and only needs process_vision_info, so the test doesn't pull transformers into the CPU image.