fix(train): repair --resume for ordinary models and fine-tuning runs; add SFT workflow - #137
Merged
Conversation
Two resume bugs found by feeding the new SFT datasets through the trainer for the first time. Both are regressions on main; the binary currently deployed on the training VM predates them and resumes correctly. Muon parameter selection was captured from the freshly built model, but load_training_state then calls restore_parameter_ids and replaces the model's ParamIds with the checkpoint's. Only the memory-hierarchy branch re-derived the selection afterwards, so resuming any ordinary model extracted Muon gradients under stale IDs, got none of them, and died on the first optimizer step with "Muon gradient is missing for parameter ...". This reproduces with plain causal_lm on the repo's own ordinary test model, so it would have hit the 300M SFT run on its first preemption. Existing tests missed it because they recompute IDs from the current model, and the QAT resume test resumes an already-complete run that takes no further steps. Separately, --checkpoint was declared conflicts_with --resume, yet the initial checkpoint's digest is part of the run signature. A resumed invocation therefore computed a signature with initial_checkpoint: None and failed with an unactionable configuration-mismatch error, making any run started from a pretrained checkpoint unresumable -- i.e. the whole SFT run. The flags may now be combined (weights still load from --output; the path only supplies identity) and the mismatch message names what differs. Adds workflow.sft.example.json: a validated 9-phase post-training workflow over the three new SFT corpora with foundation replay expressed as interleaved continued_pretrain phases at each pretraining stage's own geometry, since WorkflowV2 has no replay-fraction primitive. Every microbatch is 16,384 tokens to match the pretraining envelope. Verified: 585 tests including two new end-to-end regressions that reproduce both failures before the fix; clippy -D warnings and fmt clean.
…tion Without this there was no way to measure whether post-training improved summarization, instruction following, QA, or planning -- eval covered only causal LM and contrastive retrieval, so SFT would have been judged by reading generations, which this run has repeatedly shown to be misleading. eval now scores summarization, instruction_tuning, qa_reasoning, and retrieval_planning with token-weighted target-only cross-entropy and perplexity, reusing the trainer's own objective_loss so the number is definitionally the training loss without gradients. Unlike training, an over-long record is skipped and counted in `oversized_records` rather than aborting: training should refuse data it cannot encode, but an eval that dies on one held-out record is useless. Also repoints the SFT workflow at regenerated v2 corpora. The RAG QA generator emitted the question AFTER its passages, and the supervised encoder truncates the source from its tail, so at sequence_length 2048 the question was cut entirely from 25,282 of 74,709 records (33.8%) -- training the model to answer a question it could not see. v2 puts the question first: 0% lost, measured with the exact tokenizer. Instruction data was length-filtered to be abort-free at 2048 instead of needing 8192 to avoid a hard error, since nothing in those records is truncatable. Both phases keep 16,384-token microbatches while covering 4x and 2x more examples respectively. Pre-SFT baseline on step-249000 recorded in .context/eval-sft-baseline/: summarization ppl 22.5, instruction 12.3, qa-mixture 39.4, qa-rag 111.9, planning 49.7. Verified: 592 tests, clippy -D warnings, fmt clean.
…vice RNG training_decreases_loss_and_checkpoint_roundtrips configured dropout 0.1 and then compared an in-process model against the same model restored from its checkpoint after both took one training step. Dropout masks are drawn from the backend's global device RNG, which every concurrently running test that builds a model also draws from, so the two models saw different masks and their forward outputs diverged -- observed as "checkpoint max diff: 0.2531437", failing roughly one run in three locally and once on CI. Adding the supervised-objective eval tests raised the number of concurrent model builders and made a latent race actual; it is not a checkpoint or optimizer defect. Dropout is now zero, with a comment saying not to restore it: nothing in this test asserts dropout behaviour, and the comparison itself already runs under .valid() with dropout disabled. Verified: 6 consecutive full-suite runs clean, where the previous code failed 1 in 3.
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.
Found by feeding the newly built SFT corpora through the trainer for the first time. Both are regressions on
main— the binary currently deployed on the training VM predates them and resumes correctly, verified by inspecting the deployed source. The live spot-instance run is therefore not at risk.Bug 1 —
--resumebroken for every model without a memory hierarchyMuon parameter selection is captured from the freshly built model, but
load_training_statethen callsrestore_parameter_ids, replacing the model's ParamIds with the checkpoint's. Only the memory-hierarchy branch re-derived the selection afterwards. A resumed ordinary run therefore extracted Muon gradients under stale IDs, got zero of them, and died on the first optimizer step:Reproduces with plain
causal_lmon the repo's own ordinary test model — it would have hit the 300M SFT run on its first preemption. Existing tests missed it because they recompute IDs from the current model, and the QAT resume test resumes an already-complete run that takes no further steps.Bug 2 — a run started from
--checkpointcould not be resumed at all--checkpointwasconflicts_with = "resume", yet the initial checkpoint's SHA-256 is part of the run signature. A resumed invocation computed a signature withinitial_checkpoint: Noneand failed with an unactionablecheckpoint workflow or training configuration differs from this invocation. Since SFT starts from the pretrained checkpoint, the entire SFT run would have been unresumable. The flags may now be combined (weights still load from--output; the path only supplies identity), and the mismatch message names what differs.Also:
workflow.sft.example.jsonA validated 9-phase post-training workflow over the three new corpora (202k summarization pairs, 307k instruction/QA mixture, 126k RAG-task records): 842M compute tokens, 9,400 steps, every microbatch exactly 16,384 tokens to match the pretraining envelope.
Foundation replay has no framework primitive — WorkflowV2 phases take exactly one data path, so replay is expressed as four interleaved
continued_pretrainphases (17.1% of tokens) at each pretraining stage's own geometry, making their loss directly comparable to the live metrics stream.Tests
Two new end-to-end regressions that fail before the fix:
preempted_wake_run_resumes_and_keeps_optimizing(SIGKILL mid-run, resume, keep optimizing) andfinetune_from_checkpoint_resumes_only_with_the_same_initial_checkpoint.585 tests pass;
clippy -D warningsandfmtclean.