fix(train): let SFT retrieval replay run one epoch instead of a guessed step count - #140
Merged
Conversation
…ed step count
The retrieval replay phases added in the previous commit each requested 1200
optimizer steps. That number was invented, not measured, and the smallest
shard cannot supply it: the language-foundations retrieval split yields 46
steps at seq 512 / batch 8 / accumulation 4. The live SFT run completed its
summarization phase, entered retrieval replay, ran 19 steps, and died with
workflow phase `replay-retrieval-language-foundations-seq512` requested
1200 optimizer steps, but its data and epochs produced 46
then looped 34 times under the supervisor, restoring the same checkpoint and
failing identically for about four hours. validate-workflow cannot catch this:
it checks configuration, and the shortfall is only visible by reading the data.
`steps` is optional, so the four replay phases now run one epoch of whatever
their shard holds. This costs a one-time startup scan -- the trainer must
enumerate each epoch-bounded phase's data to compute its length, measured at
about six minutes for this workflow -- in exchange for phases that adapt to
the corpus instead of asserting a size it may not have.
Changing the workflow changes the run signature, so the interrupted run's
step-1500 checkpoint could not be resumed; it is preserved locally and in GCS
and the run restarted from the pretrained weights.
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
1 participant
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.
Follow-up to #138, found by the live SFT run crashing.
What happened
The four retrieval replay phases I added in #138 each requested 1,200 optimizer steps. That number was invented rather than measured, and the smallest shard cannot supply it — the language-foundations retrieval split yields 46 steps at seq 512 / batch 8 / accumulation 4.
The live run completed its summarization phase (1,500 steps, loss 3.12 → 2.30), entered retrieval replay, ran 19 healthy steps at 38k tok/s, then died:
It then looped 34 times under the supervisor — restore checkpoint 1500, run ~20 steps, die — for about four hours. The supervisor behaved correctly; it cannot distinguish a transient failure from a deterministic one.
validate-workflowcannot catch this. It validates configuration; the shortfall is only visible by reading the data.Fix
stepsis optional in WorkflowV2, so the replay phases now run one epoch of whatever their shard actually holds.The trade-off is real and worth stating: epoch-bounded phases force the trainer to enumerate each phase's data at startup to compute its length. Measured at ~6 minutes for this workflow, with the GPU idle. The first run started training in seconds precisely because every phase had a hardcoded step count. I consider that a good trade — phases that adapt to the corpus beat phases that assert a size the corpus may not have — but if startup cost matters more than adaptability, the alternative is explicit per-shard step counts measured against real capacity.
Cost
Changing the workflow changes the run signature, so the interrupted run's step-1500 checkpoint could not be resumed. It is preserved (
checkpoint-abandoned-1500on the box, and the 17.6 GB GCS prefix moved aside rather than deleted) and the run restarted from the pretrained weights — about 1.6 hours of A100 time.Verified: workflow validates locally and on the training host; the restarted run is training at 89% GPU.