Fix epoch indexing for experimental schedules - #6811
Open
guptaishaan wants to merge 1 commit into
Open
Conversation
Convert TrainerState.epoch to an integer before indexing list-valued beta, alpha, and mixture coefficient schedules. Treat an unset epoch as epoch zero so the properties are also safe before training starts. Add regression coverage for all three experimental trainers.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
The schedule properties indexed Python lists with
TrainerState.epoch, which isNonebefore training and a float during training. This converts the value to an integer epoch, defaulting to epoch zero, in OnlineDPO, XPO, and Nash-MD.Regression tests cover pre-training access, fractional epochs, epoch boundaries, and schedules shorter than the training run. The three affected test modules passed on Linux using CPU execution, with 41 passed and 13 skipped. Ruff lint and format checks passed.
GPU execution was not verified because the installed PyTorch CUDA runtime required a newer driver. The combined pre-commit entry point could not run because the node's Git lacks
git ls-files --deduplicate; its configured Ruff checks passed directly.Thanks to @mmjerge for reporting the issue and identifying the affected trainers.
Fixes #6806
Before submitting
AI writing disclosure
Who can review?
Anyone in the community is free to review the PR once the tests have passed.
Note
Low Risk
Narrow bugfix to epoch indexing for optional list schedules in three experimental trainers; behavior change is intentional (fixes crashes/wrong indices) with regression tests.
Overview
Fixes list-based hyperparameter schedules in Online DPO, XPO, and Nash-MD when
TrainerState.epochisNone(before training) or a float during training.Schedule properties (
beta,alpha,mixture_coef) now useint(self.state.epoch or 0)instead of indexing withself.state.epochdirectly, so pre-training access uses epoch 0, fractional epochs truncate to the current integer epoch, and out-of-range epochs still clamp to the last schedule entry.Adds parametrized unit tests in the three experimental trainer test modules for
None, fractional, and boundary epoch values.Reviewed by Cursor Bugbot for commit 63cb3ac. Bugbot is set up for automated code reviews on this repo. Configure here.