Repository navigation
[bugfix] restore LR scheduler progress on resume - #664
Conversation
Checkpoints omitted LR scheduler state, so resumed training restarted the schedule even when optimizer state was restored. Save per-rank scheduler state and restore it with an opt-in CLI flag, reconstruct legacy progress when needed, and persist epoch advancement before checkpointing.
WhiteSwan1
changed the title
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| if isinstance(scheduler, ConstantLR): | ||
| count = 0 | ||
| elif scheduler.by_epoch: | ||
| count = dataloader_state.get(EPOCHS_COMPLETED) |
There was a problem hiding this comment.
The "dataloader-state file exists but lacks EPOCHS_COMPLETED" combination is treated as unknowable progress here and hard-errors — but that combination is unambiguous: the key is only added by the StopIteration handler in main.py when the first pass completes, and neither update_dataloder_state nor the per-source consumers ever add or remove it (per the comment at line 1263). So file-present + key-missing ⇒ the checkpoint was saved mid-epoch-0 ⇒ the correct count is 0.
This affects realistic flows: epoch-mode runs that checkpoint by steps or event-time (save_checkpoints_epochs=0), and fine-tune runs (main.py pops EPOCHS_COMPLETED at startup, so their first-pass checkpoints lack the key). Resuming those from a legacy checkpoint with --restore_lr_scheduler errors out and tells the user to omit the flag — which silently restarts the schedule, exactly the harm this PR prevents.
Distinguishing "no file" (genuinely unknowable → error) from "key missing" (→ 0) would make mid-epoch-0 checkpoints reconstructible:
dataloader_state = restore_dataloader_state(checkpoint_dir)
...
elif scheduler.by_epoch:
count = (
dataloader_state.get(EPOCHS_COMPLETED, 0)
if dataloader_state is not None
else None
)The existing test_legacy_missing_progress cases would still pass (an empty test dir has no state file).
| world_size = dist.get_world_size() if dist.is_initialized() else 1 | ||
| rank = dist.get_rank() if dist.is_initialized() else 0 | ||
| if len(states) != world_size: | ||
| raise ValueError("Restored LR scheduler world size does not match.") |
There was a problem hiding this comment.
This makes a world-size change a hard failure whenever --restore_lr_scheduler is passed, yet the surrounding restore path explicitly supports resharded resume (restore_model carries MCH value-aware redistribution / _needs_mch_redistribution for exactly that). The rank-indexed JSON is the only world-size-dependent piece, and the legacy-reconstruction branch below is world-size-independent and would produce correct progress in this same case (under the "LR config unchanged" caveat it already warns about).
Consider falling back to the reconstruction path (with its warning) when len(states) != world_size, or at minimum extending this message the way the legacy error does ("Omit --restore_lr_scheduler to ...") so users on elastic/resharded resumes know the escape hatch.
Also worth noting: of the three advertised guards (world size, types/order, param-group counts), this is the only one without a test — a single-process case that writes a two-rank JSON and asserts ValueError, "world size" would cover it cheaply.
| path = os.path.join(checkpoint_dir, LR_SCHEDULER_CKPT_FILENAME) | ||
| if os.path.exists(path): | ||
| with open(path) as f: | ||
| states = json.load(f) |
There was a problem hiding this comment.
The function raises clean ValueErrors for world-size and type mismatches, but a corrupt file escapes as a bare exception: a zero-length/truncated file → json.JSONDecodeError that doesn't name the path, and a non-list top level → TypeError/KeyError at lines 1307-1308. A truncated file is reachable without a process-level crash: there's no f.flush()/os.fsync() before os.replace, and latest_checkpoint doesn't consult _SUCCESS, so a node crash after the rename persisted but before the data did leaves a resume-blocking file that must be deleted manually — and "resume after a crashed run" is exactly where an operator needs an actionable error.
An isinstance(states, list) guard plus wrapping the parse in a ValueError naming the path would match the care taken two lines below. Same theme, optional: BaseLR.load_state_dict (lr_scheduler.py:54) raises a bare KeyError when base_lrs/_last_lr are missing; a .get-style check raising the same descriptive ValueError would be consistent with the length validation there.
| - --edit_config_json: 命令行以json的方式动态修改配置文件，如{"model_dir":"experiments/","feature_configs[0].raw_feature.boundaries":[4,5,6,7]} | ||
| - --ignore_restore_optimizer: 加载checkpoint参数时，忽略加载优化器的参数 | ||
| - --restore_lr_scheduler: 从选中的checkpoint恢复学习率调度器，默认false，保持旧版本从头调度的行为。该参数独立于`--ignore_restore_optimizer`，后者仍只控制优化器状态的恢复。 | ||
| - 续训时使用`--continue_train --restore_lr_scheduler`，恢复稀疏、稠密和part optimizer的调度进度及当前学习率；续训第一批即使用恢复后的学习率。新checkpoint会保存调度状态，恢复时要求world size、调度器类型/顺序及各rank参数组数量一致。 |
There was a problem hiding this comment.
“新checkpoint会保存调度状态” sits inside the sub-bullet describing --continue_train --restore_lr_scheduler, which reads as if saving were gated on the flag. In the code, main.py passes lr_schedulers=lr_scheduler to every maybe_save/save call unconditionally (main.py:654, 688, 731), so all new checkpoints carry lr_scheduler.json regardless of the flag; only loading is opt-in (the third sub-bullet implies this but never states the save side).
Suggest making it explicit, e.g. “无论是否传入该参数，新版本保存的checkpoint都会写入调度状态（lr_scheduler.json）；不传此参数时不会加载。” Relatedly, the “Checkpoint信息与完成标记” section below documents the checkpoint dir contents for external consumers (meta, _SUCCESS) — one line naming lr_scheduler.json there would keep the layout description complete.
| None, | ||
| train_config.fine_tune_ckpt_param_map, | ||
| dense_ema=dense_ema, | ||
| lr_schedulers=lr_scheduler if restore_lr_scheduler else None, |
There was a problem hiding this comment.
Fine-tune launches (--fine_tune_checkpoint, restore_from_model_dir=False) go through this same restore, so with the flag set the new job inherits the source model's schedule progress: the modern branch restores the source's full state (base_lrs, step counts), and the legacy branch reconstructs from the source's meta["step"] / dataloader_state.json — the latter being the exact file main.py:859-861 deliberately pops EPOCHS_COMPLETED from because "fine-tune checkpoints do not carry this job's epoch budget". The scheduler path re-reads it from disk and bypasses that pop.
Net effect: fine-tuning a source that trained 10 epochs with num_epochs=2 silently trains at ~base_lr * factor^10 from batch 0, while this job's step/epoch budget starts fresh — the budget is reset but the schedule keyed to it isn't. Pre-PR, fine-tuning always got a fresh schedule, and the docs frame this flag for 续训 (--continue_train --restore_lr_scheduler).
Suggest skipping scheduler restore when not restore_from_model_dir (e.g. lr_schedulers=lr_scheduler if restore_lr_scheduler and restore_from_model_dir else None at both call sites), or at minimum warning when the flag is combined with a fine-tune checkpoint.
| use_dense_ema: bool = False, | ||
| lr_schedulers: Optional[List[BaseLR]] = None, | ||
| ) -> None: | ||
| """Restore model/optimizer state from a checkpoint dir.""" |
There was a problem hiding this comment.
Nit: this docstring predates the lr_schedulers parameter. restore now optionally restores scheduler state too, and can raise ValueError (world-size / types-order mismatch, unreconstructable legacy progress) — a new failure mode for a public entry point, worth a line like the one maybe_save got for its new arg. It would also be a good place to note that a state-based restore makes the checkpoint's schedule config (base_lrs, decay/warmup hyperparameters, even _by_epoch) win over the current config — consistent with how optimizer-state restore behaves, but currently only the legacy branch hints at it via its "LR configuration must be unchanged" warning.
|
|
||
|
|
||
| def save_lr_schedulers(checkpoint_dir: str, schedulers: List[BaseLR]) -> None: | ||
| """Save per-rank scheduler states, including rank-local parameter groups. |
There was a problem hiding this comment.
Two small docstring notes:
- The function contains a collective (
dist.all_gather_objecton the default group) but doesn't state the lockstep invariant the repo documents elsewhere (save_dataloader_state's "aggregating from all ranks first",maybe_save's lockstep comments). All current call sites are rank-uniform — but a future caller doing a rank-0-onlyCheckpointManager.save(..., lr_schedulers=...)would hang every other rank in the gather until the NCCL timeout. One sentence is cheap insurance. - "including rank-local parameter groups" slightly overstates the contents: the file stores each scheduler's
state_dict()per rank (the per-group LR listsbase_lrs/_last_lr), not parameter groups; on restore only the group count is validated while the live groups keep their params. Something like "including each rank's per-group learning rates" would be precise.
| ) | ||
| for scheduler, count in zip(schedulers, counts): | ||
| scheduler.set_step(count) | ||
| logger.info(f"Restored LR schedulers from {checkpoint_dir}.") |
There was a problem hiding this comment.
Minor: this logger.info and the logger.warning at line 1341 run on every rank, while the adjacent helpers gate informational logging to rank 0 / local-rank 0 (save_dataloader_state line 1433, restore_dataloader_state lines 1460-1462). On a large job every restore prints N copies of both lines. Consider the same int(os.environ.get("LOCAL_RANK", 0)) == 0 gate restore_dataloader_state uses. (The ValueError raise paths should stay ungated — failing uniformly on all ranks is exactly right there.)
| self.assertAlmostEqual(scheduler.optimizer.param_groups[0]["lr"], 0.0025) | ||
|
|
||
| @parameterized.expand([(False,), (True,)], name_func=parameterized_name_func) | ||
| def test_legacy_missing_progress(self, by_epoch): |
There was a problem hiding this comment.
Coverage gap worth one more case: test_legacy_missing_progress_1 (by_epoch=True) uses the bare test dir, which has neither meta nor a model.ckpt-N name — so step is also None, and the test can't distinguish "epoch schedules require EPOCHS_COMPLETED" from "nothing is known at all". If the by_epoch branch regressed to fall back to step + 1 (silently reinterpreting a batch index as an epoch count → badly wrong LRs), every current test would still pass.
Suggested case: a model.ckpt-3 dir (or meta with step) + by_epoch scheduler + no EPOCHS_COMPLETED → still raises "Cannot reconstruct".
Optionally, in lr_scheduler_test.py: test_restore_progress round-trips only LinearDecayLR; parameterizing the scheduler class — most divergently CosineAnnealingWarmRestartsLR with T_mult > 1 (cycle position via math.log) and ManualStepLR with warmup=True (bisect branch) — would cheaply pin the "n real steps ≡ set_step(n)" promise the legacy path relies on for arbitrary user configs.
| writer.close.assert_called_once_with() | ||
|
|
||
|
|
||
| class TrainLRSchedulerResumeTest(unittest.TestCase): |
There was a problem hiding this comment.
Nit (readability): this class carries the PR's core resume semantics, but the reference-replay arithmetic is opaque — expected = reference[4:] (line 707), reference[step + 1 :] (line 735), [lr * 2 for lr in expected], and pass_sizes = 3 - (skip_steps + 1) % 3 each encode non-obvious reasoning (ckpt-3 ⇔ 4 completed batch advances; being one epoch behind with decay_factor=0.5 ⇔ exactly 2× LRs; the pass-size formula replays only the remaining batches of the interrupted pass). Unlike the sibling classes here (MainTest, TrainStepCounterMultiPassTest) and in checkpoint_util_test.py, it also has no class docstring.
Tests are ruff-exempt, but a two-three line class docstring stating the invariant — "a resumed run must replay the reference LR sequence from the checkpoint onward; the default (no restore) must restart it from the beginning" — plus a short comment on the 4 and the 2×, would make a future failure here diagnosable instead of mysterious. Related: test_cold_start_with_restore_enabled can't fail on logic today (with ckpt_path=None the flag is never read — both runs execute identical code); fine to keep as a smoke guard, but the docstring could say that's its role.
Multi-agent static review — LR scheduler checkpoint/restoreFive specialized passes (code quality, performance/distributed, test coverage, documentation accuracy, security/multi-process) all completed. Overall: a solid, well-designed change with unusually thorough tests. The core semantics survived independent tracing; the findings posted inline are edge-case hardening rather than mainline bugs. Verified independently (no action needed):
Key findings (details inline):
Praise: the two-rank Note: the PR is still marked 🤖 Generated with Claude Code |
A rank keeps one fused-optimizer parameter group per embedding shard it hosts, so a restart whose sharding plan differs from the one that wrote the checkpoint restores scheduler state of the wrong width and training stopped with "Restored LR scheduler parameter groups do not match". These schedulers derive their rate from the step count alone, so the restore now falls back to the step-based rebuild that legacy checkpoints already use, and logs both group counts. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Qva5nKq3aMSc4T8oVBiyQE (cherry picked from commit bfd6983)
The cherry-picked fix carried two E501 violations. Hoist the two group counts into locals so the warning fits, and let ruff format wrap the assertion in the new test. The logged message is unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two ways a resumed run could end up on a schedule nobody asked for. torch's LRScheduler.state_dict() returns the whole instance __dict__, so the checkpoint carried the schedule's configuration next to its progress and load_state_dict's __dict__.update reinstated it. An edited pipeline.config was discarded without a word: a lowered learning_rate, a widened num_training_steps, a changed warmup. _by_epoch came back too, and that one decides whether the trainer advances the scheduler per batch or per epoch, so a checkpoint could silently change the cadence of the whole resumed run. Restore the position alone and replay it onto the live configuration, logging any field that drifted. Position is also all a rank needs, so the restore no longer depends on parameter-group counts matching and a replanned restart resumes on its own shards. --restore_lr_scheduler also fired on --fine_tune_checkpoint, importing another job's schedule position: a checkpoint at batch 100000 pinned a 5000-step fine-tune at min_learning_rate for its whole run. train_and_evaluate already drops that checkpoint's step and pops its epoch budget; gate the schedulers the same way and say so in the log. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Unit Test CPU CI runs train_eval with USE_FSSPEC=1 and
--model_dir file://..., and save_lr_schedulers died on every checkpoint:
FileNotFoundError: 'file://./tmp/.../model.ckpt-7/lr_scheduler.json.tmp'
-> 'file://./tmp/.../model.ckpt-7/lr_scheduler.json'
The file was written. filesystem_util.apply_monkeypatch() replaces open()
and os.makedirs() with fsspec-backed versions, so the .tmp landed at the
resolved local path with the scheme stripped. os.replace() is not in that
patch set, so it took the raw string and looked for a directory literally
named "file:". Nothing else under model_dir renames: save_dataloader_state
and save_meta open the final path and write it.
Do the same here and drop the temp file. That trades atomicity on the
epoch-boundary refresh for portability, so restore now treats an
unreadable file as absent and rebuilds from the checkpoint step rather
than ending the run. Covered by a round trip through a file:// URI with
the monkeypatch applied, and by a truncated-payload test.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…recorded A checkpoint with no lr_scheduler.json was resumed by inferring the schedule position from its step counter. That inference is unsound. --restore_lr_scheduler defaults false, but the save at main.py:668 is unconditional, and a resume with the flag off skips replayed batches (main.py:617-618) before the per-batch lr.step() (main.py:639-641) -- so the schedule deliberately lags the step counter and the checkpoint that run writes records a position well below step + 1. Rebuilding from the step counter on that lineage is not a restore, it is a guess, and a two-generation run measures it at 5x the recorded position with only a warning to show for it. So stop guessing. The file records one integer per scheduler per rank and that integer is the whole feature; absent it there is nothing to restore, which is the ordinary state of a checkpoint written before this existed, not an error. Warn, start the configured schedule from the beginning, and say which flag to drop. A file that exists but will not parse is the opposite case -- a position was recorded and is gone -- so that raises, naming the file and the two ways out. Removes _set_lr_schedulers_from_progress, the ConstantLR carve-out that existed only to spare it a meaningless epoch count, the dependency on meta / __epochs_completed__, and both fallback handlers. by_epoch resume no longer needs a completed-epoch count at all. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The wrapper bought one thing the raw error lacks -- the path, since json.load() takes a handle -- and paid for it by discarding the JSONDecodeError subclass, at eight lines, in a function whose every other error is a plain raise and whose neighbour restore_dataloader_state lets json.load fail untouched. The traceback names the frame, the checkpoint is logged just before it, and train.md already documents the remedy, so the path is recoverable without inventing an exception type. Behaviour is unchanged: a recorded-but-damaged state still ends the run rather than quietly starting the schedule over. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The checkpoint stored scheduler.state_dict(), which torch defines as the
whole instance __dict__ minus the optimizer. So the on-disk shape was
pinned to a torch implementation detail that can change without anyone
editing this repo, 12 of its 13 fields were decorative, and the one field
a restore needs sat among torch's bookkeeping. _PROGRESS_KEYS existed to
subtract that bookkeeping back out -- a denylist over another library's
private attributes, which every torch bump could silently change with no
test to notice.
Override state_dict() instead: {position, config}. BaseLR.__init__
snapshots the attributes already on the instance, which is exactly the
subclass's configuration since every subclass assigns before delegating
to super, so the allowlist is captured rather than maintained and a new
scheduler class needs no bookkeeping. _step_count, _last_lr, _is_initial
and _get_lr_called_within_step stop being written; set_step recomputes
them. _last_lr was one float per param group, so a sharded job halves its
per-group data.
This also restores the reader/writer symmetry: load_state_dict already
had non-torch semantics while state_dict() was inherited unchanged, so
load_state_dict(state_dict()) no longer meant what torch's contract says.
The drift warning now diffs a declared config rather than a subtraction.
Carries the review cleanups from the same pass: maybe_save's Args block
still promised a per-rank save, the save/restore docstrings narrated the
deletion instead of the code and were the only two in their group without
Args, the types/order mismatch now prints expected vs actual, the corrupt
-state test fixture still spelled the retired nested layout, and two test
names described designs that no longer exist.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The checkpoint recorded each scheduler's configuration alongside its position, for one purpose: warning when that configuration differed from the resumed run's. But the live configuration always wins -- that is the contract, and it is documented -- so the warning fired on the feature working as designed. Lowering learning_rate for a second phase earned a line per scheduler telling you the change had taken effect. Record only the position. The configuration is rebuilt from pipeline.config on every launch, so storing it only created something a restore could wrongly prefer over the live values. Dropping it takes the _config_keys capture, the drift comparison, the warning and this module's logger import with it, and turns a behavioural guarantee into a structural one: nothing else is recorded, so nothing else can be adopted, and the silent-config-clobber bug is no longer expressible. On a two-rank sharded job the file goes from 255 bytes to 47. No test covered the warning, so none had to change -- the test that pins this behaviour asserts the configuration values survive a restore, which is now guaranteed by construction rather than by load_state_dict choosing not to apply them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
459 test lines against 167 of production invited a cull, but most of them earn their keep: of twenty candidates, seventeen survived an adversarial pass that hunted for a production edit each one alone would catch, and several were confirmed by mutation sweeps. The parameter tuples are not padding -- they pin distinct lines across the step, epoch, dedupe, fine-tune and resharding paths. So this is mostly abstraction rather than deletion. Both scheduler fixtures collapse into one module-level helper, placed beside the other module helpers so the two-rank worker can reach it, and every LR checkpoint test now lives in LRSchedulerCheckpointTest instead of being split across two classes that each built their schedulers a different way. lr_scheduler_test grows one _sgd() that replaces four spellings of the same optimizer construction. Three tests go. test_restore_lr_schedulers_uses_saved_state asserted a state_dict round trip that test_round_trip already covers. epoch_resume's mid_epoch tuple exercises the same path as epoch_boundary_step_save. test_fine_tune_ignores_restore_lr_scheduler is merged, not dropped -- its coverage of the restore_from_model_dir gate is real and unique, so it becomes a sixth batch_resume tuple, assertion intact. Shrinks: test_restore_keeps_this_run_configuration no longer builds a third scheduler just to compute an expected rate, and moves to position 12, where the two schedules actually differ -- at position 4 it asserted a value both produced. The by_epoch axis leaves the no-saved-state test, which stopped depending on it when the reconstruction path was deleted. 459 -> 409 test lines; 131 tests, all passing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The entry had grown to seven sub-bullets, most of them documenting behaviour that no longer exists: the step-based reconstruction, the drift-warning field list, the corrupt-file recovery, the per-rank layout and its world-size requirement. Every neighbouring flag gets one line. Keep the four things the flag name does not already say: what it restores, that it defaults to off, that it applies only to --continue_train, and that pipeline.config still decides the rates. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
WhiteSwan1
changed the title
|
Automated review of 2026-09-22 at head Addressed
Not changed
No longer applicable
Notes for re-reviewTwo statements under Verified independently no longer describe the code: the Renamed since the review: The branch conflicts with base |
test_restore_progress and its legacy parameter were vocabulary from the step-based reconstruction, which no longer exists; the code says position. Rename to test_restore_position / via_set_step, which also says what the axis is for: load_state_dict is a thin wrapper over set_step, and a clamp added inside it survives the set_step branch while killing the other. TrainLRSchedulerResumeTest gains the class docstring the review asked for, stating the invariant every case compares against, plus one line each on the 4, the 2x and the pass_sizes formula. The rest of the comment volume added during consolidation is cut roughly in half: what remains explains arithmetic, not what the line below it already says. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
| - --fine_tune_checkpoint: 增量训练的checkpoint路径，如experiments/multi_tower_din_taobao_local/model.ckpt-0，如果不设置，增量训练使用model_dir下最近的检查点 | ||
| - --edit_config_json: 命令行以json的方式动态修改配置文件，如{"model_dir":"experiments/","feature_configs[0].raw_feature.boundaries":[4,5,6,7]} | ||
| - --ignore_restore_optimizer: 加载checkpoint参数时，忽略加载优化器的参数 | ||
| - --restore_lr_scheduler: 续训时从checkpoint恢复学习率调度进度，默认false即从头调度；只对`--continue_train`生效，且只恢复进度，学习率配置仍以`pipeline.config`为准 |
There was a problem hiding this comment.
This bullet invites users to edit the LR config and resume ("学习率配置仍以pipeline.config为准"), but three resume behaviors are worth documenting here:
- Editing schedule values takes effect, but changing the scheduler type/count/order between save and resume makes
restore_lr_schedulersraiseValueError("Restored LR scheduler types/order do not match...")and training fails at startup (tzrec/utils/checkpoint_util.py:1271-1275). Users switching LR strategy type must drop--restore_lr_scheduler(schedule restarts fresh). --continue_trainis necessary but not sufficient: combined with--fine_tune_checkpointand a freshmodel_dir, the restore comes from the fine-tune checkpoint and the flag is ignored (only a runtime warning,tzrec/main.py:580-585). One clause saying the flag never applies to fine-tune checkpoints would prevent the surprise.- The most common case right after release — resuming a checkpoint written before this feature — hits the "no scheduler state" warning path and restarts the schedule from the beginning (
checkpoint_util.py:1261-1266). A short 旧版checkpoint无调度状态文件时会warning并从头调度 clause covers the troubleshooting case the docs convention asks for.
| def state_dict(self) -> Dict[str, Any]: | ||
| """Persist the schedule position, and nothing else. | ||
|
|
||
| torch would return the whole instance ``__dict__``, pinning the |
There was a problem hiding this comment.
Nit on accuracy: torch's LRScheduler.state_dict() returns every entry of self.__dict__ except optimizer, not "the whole instance __dict__". The rest of the rationale holds; something like "nearly the whole instance __dict__ (everything but the optimizer)" would be exact.
| # written in place like the dataloader state beside it: model_dir may be an | ||
| # fsspec URI, where only the patched helpers in filesystem_util work and | ||
| # os.replace does not, so an epoch-boundary refresh overwrites, not renames. | ||
| with open(path, "w") as f: |
There was a problem hiding this comment.
Judgement call, flagging for visibility rather than asking for a change: this in-place truncating write is the one write that can damage an already-published checkpoint. The dedupe refresh in maybe_save (line 576) rewrites lr_scheduler.json inside the _SUCCESS-marked _last_ckpt_dir, and latest_checkpoint does not filter on _SUCCESS. If rank 0 is killed between truncate and flush — the preemption window this flag is meant to serve — every subsequent --restore_lr_scheduler restart dies on json.JSONDecodeError until an operator deletes the file or drops the flag, even though the rest of that checkpoint is intact.
The window is tiny (~100 bytes) and this faithfully mirrors the pre-existing dataloader_state.json refresh beside it, so accepting it for parity is reasonable; just making sure it is a conscious choice. If you ever want to close it, local-path write-temp + os.replace handles the non-fsspec case without touching the URI path (reader-side tolerance would conflict with the deliberate "damaged is an error" stance pinned by test_unreadable_state_is_an_error).
| """Replay the recorded position onto the schedule this run configures. | ||
|
|
||
| A learning rate edited in pipeline.config takes effect on resume, and | ||
| so does ``_by_epoch`` -- which decides whether the trainer advances |
There was a problem hiding this comment.
The recorded position is unit-blind: its unit (batch vs epoch advances) comes from the live config's by_epoch, which is deliberately not recorded, and the on-disk check only compares class names. Flipping by_epoch in pipeline.config (a plain proto field, default false) across a resume therefore passes validation and replays e.g. 50k step-advances as 50k epoch-advances — LinearDecayLR/CosineAnnealingLR clamp (t = min(...)) and the entire resumed run silently trains at min_learning_rate; ExponentialDecayLR decays toward zero. The same silent clamp hits without any flip if the resume config shrinks num_training_steps/T_max below the recorded position.
I see this is a chosen trade-off (test_restore_keeps_this_run_configuration pins config-wins-over-units), but today a mis-edited resume degenerates the LR with no signal in the logs. A minimal mitigation that keeps the pinned behavior: record "by_epoch" next to "position" in state_dict (the on-disk format is new in this PR, so it is free to shape) and logger.warning — not raise — on mismatch. At minimum, this docstring currently sells the flip as unqualifiedly safe; it should mention that a position recorded in one unit is replayed in the other.
| something a restore could wrongly prefer over the live configuration. | ||
|
|
||
| Returns: | ||
| the number of advances the schedule has made. |
There was a problem hiding this comment.
Tiny: Returns describes the value inside the dict, but the function returns {"position": int}. Since this docstring defines the on-disk contract save_lr_schedulers/load_state_dict rely on, describing the shape would be clearer, e.g. "A dict with a single key position: the number of advances the schedule has made."
| env.start() | ||
| self.addCleanup(env.stop) | ||
|
|
||
| def test_round_trip(self): |
There was a problem hiding this comment.
As written, this round trip cannot detect a save/restore order swap: both schedulers are ExponentialDecayLR at set_step(4), so the two JSON entries are byte-identical and the zip(original, resumed) assertions pass even if the list came back reversed. Production restores [sparse_lr, dense_lr, *part_lrs], which can hold same-type schedulers at different positions (sparse step-based, dense epoch-based), and test_scheduler_type_mismatch only catches reordering when types differ — the case the type check already handles. One-line fix: distinct positions (e.g. set_step(4) / set_step(7)).
| ), | ||
| mock.patch("tzrec.main._log_train"), | ||
| mock.patch("tzrec.utils.checkpoint_util.save_model"), | ||
| mock.patch("tzrec.utils.checkpoint_util.restore_model") as restore_model, |
There was a problem hiding this comment.
Nothing in the suite pins that CheckpointManager.restore applies scheduler state after restore_model's optimizer.load_state_dict — and that order is load-bearing: torch's optimizer load restores param_groups["lr"] from the checkpoint, so a flip inside restore() would silently clobber the just-restored schedule (for a by_epoch scheduler, for a whole epoch). Every test mocks restore_model, so the flip would pass the suite.
Cheap pin: give this mock a side effect that overwrites param_groups[0]["lr"] with a sentinel when an optimizer is passed; the existing reference[4:] assertions in the "resume" case then prove the scheduler restore wins. That would also make the class docstring ("the LR sequence a resumed run actually trains on") honest for the no-flag cases: a real optimizer restore hands the first resumed batch the checkpoint's saved LR until the fresh scheduler's first step() overrides it (pre-existing torch behavior, not a defect of this PR), which the current mock hides.
| # fine-tune checkpoints do not carry this job's epoch budget | ||
| dataloader_state.pop(checkpoint_util.EPOCHS_COMPLETED, None) | ||
| manager = checkpoint_util.CheckpointManager(os.path.join(self.test_dir, name)) | ||
| config = TrainConfig( |
There was a problem hiding this comment.
Every harness run sets save_checkpoints_steps=1 (or epochs=1), so the trailing maybe_save(final=True) always lands in the dedupe-refresh branch and the non-dedupe final save is never exercised. That path matters: the final save is frequently the latest checkpoint a --continue_train resume picks (e.g. save_checkpoints_steps=1000 with num_steps=1500, or SID-style runs where periodic saving is disabled and the final save is the only checkpoint). Dropping lr_schedulers= from that call site (main.py:745) would pass the whole suite while the most common resume point degrades to warn + fresh schedule. Cheap pin: one source run with save_checkpoints_steps=0, then assert its model.ckpt-<last>/lr_scheduler.json records position == num_steps.
| expected = reference[:2] | ||
| torch.testing.assert_close(actual, expected) | ||
|
|
||
| @parameterized.expand( |
There was a problem hiding this comment.
The boundary=False axis is only paired with restore=False, so both restore=True epoch cases resume from model.ckpt-2 — whose epoch-unit position comes from the boundary save / dedupe refresh. A regression that made plain in-loop step saves record the wrong epoch-unit position would pass all three cases; only a restore from mid-pass model.ckpt-4 (position 1, written by a plain step save) exposes it. Adding ("epoch_mid_pass", True, False, False) costs one table row and the existing expected = reference[step + 1:] logic already handles it.
| def setUp(self): | ||
| self.test_dir = make_test_dir() | ||
| self.addCleanup(shutil.rmtree, self.test_dir) | ||
| env = mock.patch.dict(os.environ, {"RANK": "0", "LOCAL_RANK": "0"}) |
There was a problem hiding this comment.
Two cheap unit pins are missing in this class:
- The
RANK != 0early-return insave_lr_schedulersis unasserted —RANKis pinned to 0 here, and in the gloo worker both ranks would write identical content if the gate were removed, so every assertion still passes. The gate is what keeps N ranks from racing on one JSON file (a torn write is a hardJSONDecodeErroron resume, pertest_unreadable_state_is_an_error). A one-line test (RANK=1→ no file written) would cover it. set_step's negative-step guard (lr_scheduler.py:84-85) has no test either.
Multi-agent static re-review at head
|
Documentation findings, all accurate. state_dict's docstring said torch returns "the whole instance __dict__" when it returns every attribute but the optimizer, and its Returns described the value rather than the shape. train.md's one-liner omitted three behaviours a user can hit: a changed scheduler type/order hard-fails and needs the flag dropped, the flag never applies to --fine_tune_checkpoint, and a checkpoint predating the feature warns and starts fresh. Restores epoch_mid_pass, deleted in 772f2fc on a redundancy finding that was wrong. Every other by-epoch restore resumes from model.ckpt-2, whose position the dedupe refresh writes; nothing resumed from mid-pass, whose position the ordinary per-step save writes. Confirmed by mutation: drop lr_schedulers= from the in-loop save and this case alone fails, while both boundary cases pass. round_trip now uses distinct positions, so a reversed restore fails rather than passing on byte-identical entries, and asserts the decoded file whole -- which also subsumes the separate one-entry-per-scheduler test, since comparing two type names would have passed under the old per-rank layout. Cuts, each for a reason that holds against the current payload: cold_start ran identical code either way, the flag being unreachable without a checkpoint. The checkpoint-level param-group test could only duplicate the scheduler-level one, the payload having no group information left for any checkpoint-level guard to check. via_set_step covered set_step directly, which load_state_dict already reaches. The unreadable-state test pinned json.load's own behaviour. Verified on torch 2.13.0 / torchrec 1.8.0: 124 tests pass, and test_local_fsspec_train_eval_export passes -- the end-to-end check the local container cannot run, being below master's dependency floor. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The restore block printed both flag values on every restore, including ignore_restore_optimizer, which predates this change and was never logged before. A correct invocation now logs nothing. What survives is the case worth a line: the flag was passed explicitly and silently did not apply, because a fine-tune checkpoint's schedule position belongs to the source job. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
BaseLR no longer overrides state_dict; save_lr_schedulers writes the one field a restore needs, last_epoch, under a "state" key beside the class name. load_state_dict keeps torch's __dict__ update and follows it with set_step, which realigns _step_count, _last_lr and the optimizer's lr -- torch alone leaves those at the fresh scheduler's values. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Mutation testing on the remote box: every defect the count=1 row caught was also caught by count=4, and count=12 sits past the horizon where a restored and a fresh LinearDecayLR both read lr 0.0, so a stale _last_lr or parameter-group lr could not fail it. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
… bugfix/restore-lr-scheduler
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
Thanks for the second pass at What changed after the reviewThe biggest change is how scheduler state is persisted ( Smaller: the fine-tune notice is a warning and only fires when the flag is actually being ignored ( |
The five tests that drive _train_and_evaluate against mocks each built their own TrainConfig, four of them spelling out proto defaults. They now call _train_config and pass only what differs; the patch blocks and every assertion are unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
| - --fine_tune_checkpoint: 增量训练的checkpoint路径，如experiments/multi_tower_din_taobao_local/model.ckpt-0，如果不设置，增量训练使用model_dir下最近的检查点 | ||
| - --edit_config_json: 命令行以json的方式动态修改配置文件，如{"model_dir":"experiments/","feature_configs[0].raw_feature.boundaries":[4,5,6,7]} | ||
| - --ignore_restore_optimizer: 加载checkpoint参数时，忽略加载优化器的参数 | ||
| - --restore_lr_scheduler: 续训时从checkpoint恢复学习率调度进度，默认false即从头调度。只恢复进度，学习率配置仍以`pipeline.config`为准；但调度器的类型/个数/顺序改变会直接报错，换调度策略时需去掉该参数。该参数只对`--continue_train`从`model_dir`恢复生效，`--fine_tune_checkpoint`会忽略它（仅日志提示）。旧版本checkpoint没有调度状态，会打印warning并从配置的调度起点开始 |
There was a problem hiding this comment.
Restoring the LR scheduler should be the default, since failing to restore it causes the learning-rate schedule to restart incorrectly after resuming training.
I'd suggest changing the flag to ignore_restore_lr_scheduler, with the default set to false. This would make scheduler restoration the default behavior while still allowing users to explicitly opt out when needed.
| writer.close.assert_called_once_with() | ||
|
|
||
|
|
||
| class TrainLRSchedulerResumeTest(unittest.TestCase): |
There was a problem hiding this comment.
keep it in MainTest
The LR-scheduler resume harness and its cases move from their own class into MainTest, which now owns the temp dir they write to. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Restoring the schedule position was opt-in, so a resume that did not pass --restore_lr_scheduler still restarted the learning-rate schedule -- the bug this change set out to fix. The flag becomes --ignore_restore_lr_scheduler, default false, mirroring --ignore_restore_optimizer: --continue_train now continues the schedule where the checkpoint left it unless told not to. A checkpoint without scheduler state still warns and starts fresh, a scheduler type or order mismatch still fails and now names the opt-out, and fine-tune checkpoints never restore the position, so their warning is gone. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
MainTest now creates a temp dir per test for the resume harness, so the one test that made its own with tempfile can use that instead. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…he loop harness The four lr_scheduler.json tests join CheckpointUtilTest, whose setUp already provides the temp dir; their rank pin goes, since that class pins rank only where a non-zero path is under test. MainTest's resume helper is _train_lrs, after what it returns: the learning rate each batch of one run trained with. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…checkpoint A job started from --fine_tune_checkpoint does restore the position once it resumes from its own model_dir; the sentence now names the source, not the job. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
2 participants
Summary
Training checkpoints restore model and optimizer state but do not record LR scheduler
progress, so a resumed run restarts the learning rate schedule.
Resuming with
--continue_trainnow restores the schedule position by default;--ignore_restore_lr_scheduleropts out, mirroring--ignore_restore_optimizer, which controlsoptimizer state alone and is independent of it. This changes behavior for every existing resume:
the learning rate continues where the checkpoint left it instead of restarting. Epoch-based
schedules advance before the epoch checkpoint is saved, so an epoch-boundary checkpoint records
the position the next epoch starts from; the learning-rate trajectory of an uninterrupted run is
unchanged.
Every checkpoint records the schedule position — rank 0 writes one
lr_scheduler.jsonwithone entry per scheduler,
{"type": <class name>, "state": {"last_epoch": N}}; the opt-outaffects loading only. Nothing but the position is recorded: a learning
rate, warmup, horizon or
by_epochedited inpipeline.configtakes effect on resume,because the checkpoint carries nothing that could override it. Restore is therefore
insensitive to world size and to each rank's parameter-group count, so rescaling and
resharding resume correctly; only scheduler type and order are enforced, and a mismatch fails at
startup with a message naming
--ignore_restore_lr_scheduleras the way to start the new schedulefresh.
BaseLR.load_state_dictis torch's__dict__update followed byset_step, which realigns_step_count,_last_lrand the optimizer'slr. Torch alone leaves those at the freshscheduler's values, and every
_get_lrreads_step_count, so the schedule would computefrom step 0 until the next
step(). An earlier revision instead overrodestate_dicttoreturn only the position and replaced torch's
load_state_dictoutright; both were dropped.Torch's read side is harmless, the write side only needs the follow-up, and the file's
layout belongs with
save_lr_schedulers, the function that writes it.A checkpoint written before this existed records no position. That is not an error: the run
warns and starts the configured schedule from the beginning. Progress is deliberately not
reconstructed from the checkpoint's step counter — a resume that ran without the flag advances
the step counter while holding the schedule back, so the two can differ by a wide margin, and
inferring one from the other would be a guess rather than a restore.
Restore applies only to
--continue_trainresuming frommodel_dir. A--fine_tune_checkpointcarries another job's schedule position, so it is never restored there, matching the existing
handling of that checkpoint's step counter and epoch budget.
The restore policy is a CLI option rather than a TrainConfig field: it is a launch-time
decision, keeps existing configurations unchanged, and needs no protobuf change.
Test Plan
lr_scheduler_test,checkpoint_util_test,main_test): 119 tests pass on torch 2.13.0 / torchrec 1.8.0. Resume coverage:default (restores), weights-only, opt-out with and without optimizer state, a checkpoint with
no recorded position, fine-tune; epoch resume from mid-pass and from an epoch boundary under
both step-based and epoch-based saving, including the dedupe refresh; the fsspec-URI round
trip.
onto its own base learning rates, and that stepping continues correctly afterwards.
set_step(k)verified equivalent to k realstep()calls for all six schedulersubclasses at every k in 0..17.
MainTest; every mocked training-loopTrainConfigin that module comes from one
_train_configfactory, and the four pre-existing loop tests werechanged only to call it.
pre-commit runon the changed files andpyrefly check: clean.🤖 Generated with Claude Code