[Bug] Rollout data source underfills batches spanning multiple epochs

[LOGO127]

Bug Description

RolloutDataSource.get_samples(n) only handles at most one epoch wrap. If the requested group count exceeds the remaining current epoch plus one complete dataset, it returns fewer groups than requested and leaves sample_offset beyond the dataset length.

Reproduced against current main ce92eff12ecdc81396bf41a2f94e62dd5b0aca32 with actual RolloutDataSource, Dataset and a three-row local JSONL file. This is an AI-assisted CPU data-source report, not an end-to-end training failure claim.

Steps to Reproduce

Save the following as repro_cursor.py and run from the Vime checkout with its import dependencies installed:

PYTHONPATH=. HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 python repro_cursor.py

"""Run from a Vime checkout with PYTHONPATH=. and the usual CPU import dependencies."""

import json
import platform
import tempfile
from pathlib import Path
from types import SimpleNamespace

import torch

from vime.rollout.data_source import RolloutDataSource
from vime.utils.data import Dataset

print({"python": platform.python_version(), "torch": torch.__version__})
with tempfile.TemporaryDirectory(prefix="vime-cursor-repro-") as directory:
    path = Path(directory) / "prompts.jsonl"
    path.write_text("".join(json.dumps({"text": str(i)}) + "\n" for i in range(3)))
    source = RolloutDataSource(SimpleNamespace(
        rollout_global_dataset=False, prompt_data=None,
        rollout_shuffle=False, n_samples_per_prompt=2))
    source.dataset = Dataset(str(path), tokenizer=None, processor=None,
        max_length=None, prompt_key="text")
    groups = source.get_samples(8)
    print({"requested_groups": 8, "actual_groups": len(groups),
        "prompts": [group[0].prompt for group in groups],
        "sample_offset": source.sample_offset, "epoch_id": source.epoch_id})

The fixture supplies a real Dataset after disabling automatic model loading in the constructor. No tokenizer/model download, import stubs, GPU or inference server are used.

Expected Behavior

Since the source already wraps across epochs, a request for eight groups from three prompts should return:

prompts = ['0', '1', '2', '0', '1', '2', '0', '1']
sample_offset = 2
epoch_id = 2

This matches eight successive get_samples(1) calls and preserves per-epoch shuffling when enabled.

Actual Behavior / Logs

{'python': '3.12.3', 'torch': '2.11.0+cpu'}
{'requested_groups': 8, 'actual_groups': 6, 'prompts': ['0', '1', '2', '0', '1', '2'], 'sample_offset': 5, 'epoch_id': 1}

The call returns normally; there is no exception traceback. The cursor state can also be saved/restored in this out-of-range state.

Environment

• Vime source: main SHA above, rechecked immediately before reporting
• OS: WSL Ubuntu on Windows
• Python: 3.12.3; PyTorch: 2.11.0+cpu
• CUDA/ROCm/GPU: not used
• vLLM, vllm-router, Megatron-LM: not executed by this reproduction
• No full training dependency or GPU compatibility claim

Additional Context / Proposed Scope

The rollout loop passes over_sampling_batch_size into the data source. It may compensate for a short return by fetching again; I have not demonstrated a final rollout batch-size failure, training hang, or model-quality regression. The proven defect is the data-source return/cursor contract.

An unpublished local candidate iterates epoch boundaries until the requested groups are collected, retaining existing per-epoch shuffle, deepcopy and group/index behavior. Positive requests on an empty dataset fail explicitly to avoid an infinite loop. The exact-boundary convention remains unchanged (advance/shuffle at the next positive request).

Local validation includes repeated mixed batch sizes, datasets of size 1/3/7, shuffle on/off, actual cursor torch.save/load and buffered top-up:

• Original source: 16 failed / 26 passed.
• Candidate: 42 passed.
• One failure checks the proposed empty-dataset error behavior, not an already-documented exception contract.
• Scoped Ruff/Black/isort checks passed; full pre-commit, GPU rollout and human end-to-end validation are outstanding.

Would repeated epoch wrapping be the desired behavior for small datasets/large oversampling batches, or would you prefer explicitly rejecting such requests? No PR or human-review/sign-off claim yet; I will keep any implementation aligned with your contribution requirements.

Pre-report checks performed with AI assistance: read CONTRIBUTING, README and the FAQ (including batch-size/resume sections); searched issues and PRs for sample_offset, get_samples and dataset/epoch terms; reproduced on current main. No focused duplicate was found.

[LOGO127]

Additional AI-assisted CPU evidence for the impact boundary: vime.rollout.sft_rollout.generate_rollout calls data_buffer.get_samples(args.rollout_batch_size) once and returns the resulting groups without a refill loop.

I exercised that actual function on main ce92eff12ecdc81396bf41a2f94e62dd5b0aca32, using three local JSONL user/assistant conversations, n_samples_per_prompt=1, rollout_batch_size=8, rollout_global_dataset=True, and shuffle off/on. The fixture creates a local fast WordLevel tokenizer and a synthetic ChatML-shaped template, then uses the real tokenizer/processor loaders and MultiTurnLossMaskGenerator (loss_mask_type=qwen3_5). Only module-global tokenizer/mask caches are reset between tests; those functions are not mocked.

• Unchanged main: both cases fail the batch-length assertion. SFT returns 6 groups, epoch 1 / offset 5.
• Existing unpublished data-source candidate: both cases pass. SFT returns 8 groups, epoch 2 / offset 2; prompt order matches repeated one-group fetches. Tokens and nonempty response loss masks are checked, and a second SFT batch after actual cursor save/load matches uninterrupted execution.

This establishes a short return at the SFT rollout function boundary, in addition to the data-source-only reproduction above. It does not establish a trainer crash, optimization/model-quality impact, production Qwen tokenizer compatibility, or a GPU/Ray/Megatron end-to-end run. No training or inference service was launched, and human validation remains outstanding. No PR or ownership claim is added by this follow-up.