Skip to content

fix: resume SFT dataset iteration at the exact batch - #908

Open
RaycarlLei wants to merge 1 commit into
OpenPipe:mainfrom
RaycarlLei:fix/sft-resume-batch-offset
Open

RaycarlLei wants to merge 1 commit into
OpenPipe:mainfrom
RaycarlLei:fix/sft-resume-batch-offset

Conversation

@RaycarlLei

Copy link
Copy Markdown

initial_step is a global batch offset, but create_sft_dataset_iterator converts it to a global chunk offset. This replays completed batches when an epoch ends with a short chunk, or when resuming inside a chunk.

For example, with 5 trajectories, batch_size=2, chunk_size=2, and two epochs, the chunk steps are [0, 2, 3, 5]. Resuming at step 3 currently yields [2, 3, 5]; resuming at step 6 still yields step 5 even though training is complete.

Trim each chunk's start using the batch offset within its epoch, preserving the original chunk endpoint. The existing learning-rate slicing and step metadata then start at the same batch as the remaining trajectories.

The regression test compares resumed batches against the uninterrupted stream's suffix, including trajectories, learning rates, global steps, epoch, and epoch steps. It covers mid-chunk and epoch boundaries, completion, full/partial final batches, and deterministic shuffling. On the original implementation, 24 of the 32 new cases fail.

Validation:

  • uv run pytest tests/unit/test_sft.py -q --tb=short: 47 passed on Windows and Linux (Python 3.12).
  • uv run prek run --all-files: Ruff, formatting, full src/tests type checking, and lock validation all passed on Linux.
  • Linux verification checks the exact contribution commit, 6f11634280192dd4664f634b61a681cfe30c77ac.

This validates the dataset iterator; GPU training and the backend integration suite were not run.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant