feat(train): corrective SFT workflow — right-sized replay, v3 abstention data, generative final phase - #141
Merged
Merged
Conversation
…fects The first SFT run improved every held-out objective sharply -- planning ppl 47.0 -> 3.78, qa-rag 138.1 -> 8.49, instruction 11.7 -> 5.90 -- while leaving free-running generation degenerate: summarization copies its first sentence and loops, instruction emits numbered questions instead of answers, planning echoes the request. Teacher-forced loss rewards knowing the format given a correct prefix; autoregressive decoding compounds error into copying, and perplexity cannot see it. Three changes, each measured rather than guessed: Replay steps are explicit and short. Epoch-bounding fixed the earlier crash but let phase 10 run ~10,900 steps on the 1.3GB advanced shard -- 40% of the run -- so 12,200 of the final 13,100 steps were replay. Retrieval repair completes in 17 steps (measured: phase 8 reached loss<1e-4 after 17 of 2,449), so 150 is generous. Replay is now 7% of compute instead of 63%. RAG QA reads the v3 corpus. Every unanswerable record in v2 shared one target string, making it 20,606 of 74,709 targets and the model's favourite output for any prompt. v3 uses 34 structurally varied paraphrases at a 10% share; the most frequent target fell from 27.6% to 0.48%. The run ends on a generative phase. The first run ended on causal replay, so its final gradients shaped retrieval and language modelling rather than generation. Starts from the first SFT checkpoint, which already holds the retriever (98.0/99.7/100.0% top-1, unchanged by SFT) and the improved language model (advanced ppl 121.6 -> 81.5). 10 phases, 6,550 steps, 442.6M compute tokens -- roughly a third of the first run.
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.
The first SFT run finished (26,824 steps). Held-out loss says it worked; generation says it did not. This workflow addresses why.
What the first run produced
Retrieval was completely preserved — 98.0 / 99.7 / 100.0% top-1, identical to pre-SFT — and causal improved (advanced ppl 121.6 → 81.5). The interleaved replay from #138 did its job exactly.
But free-running generation is degenerate on every objective: summarization copies sentence one and loops it, instruction emits numbered questions, planning echoes the request, grounded QA repeats the question instead of reading passage [1]. Teacher-forced perplexity rewards knowing the format given a correct prefix; autoregressive decoding compounds error into copying. Exposure bias, structurally invisible to loss.
Three fixes, each measured
Replay steps explicit and short. Epoch-bounding (#140) fixed the crash but let phase 10 run ~10,900 steps on the 1.3 GB advanced shard — 40% of the run — so 12,200 of the final 13,100 steps were replay, and the run ended on replay. Measured repair time is 17 steps (phase 8 hit loss <1e-4 after 17 of 2,449), so 150 is generous. Replay drops from 63% of compute to 7%.
RAG QA reads v3. In v2 every unanswerable record shared one target string — 20,606 of 74,709 targets, 85× the next most frequent — and the mid-run checkpoint emitted it for every prompt, even with the answer verbatim in passage [1]. v3 uses 34 structurally varied paraphrases at a 10% share; most frequent target 27.6% → 0.48%.
Ends on a generative phase, so the final gradients shape generation rather than retrieval.
Shape
10 phases, 6,550 steps, 442.6M compute tokens — about a third of the first run. Starts from the first SFT checkpoint, which already carries the retriever and the improved language model. Validates against
retriever_300m_moe.mal.Retrieval replay is retained (150 steps × 3) because the first run proved it necessary — without it the retriever would erode, as it did during pretraining's causal-only stages (42.5% → 21.3% top-1).