Repository navigation
Conversation
6997c74 to
de6fb63
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Arist12
changed the title
d651c3c to
12f7112
Compare
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
I reviewed the current head (55dc933d6c9070940f854a88f9111f018d98853c) and exercised the implementation with a real two-rank Megatron distributed Adam optimizer on AMD Instinct MI355X GPUs.
The happy path restored model-side adapter tensors, fp32 master weights, first/second moments, optimizer metadata, and iteration exactly. The following failure paths still need work.
Blocker — parameter-state structure is not validated before DP scatter.
_load_optimizer_param_state coordinates only torch.load. With a valid pickle containing an invalid state structure, DP rank 0 raised KeyError while checking numel_unpadded, while rank 1 had already entered scatter and failed after the peer disconnected. This was reproduced with the real Megatron DistributedOptimizer, not a mock.
DP roots need a non-mutating structural/numel preflight, followed by global consensus, before any rank calls load_parameter_state_from_dp_zero.
The final rebase must preserve #2733 before any mutation.
#2516 conflicts exactly at the native-load site and its side contains the permissive if name in state_dict: copy_ loop. The combined flow must be:
- every rank reads its selected native shard;
- strict key/shape/dtype/finiteness validation runs locally;
- ranks reach global error consensus;
- only then copy model parameters and reload optimizer masters;
- preflight and restore optimizer/scheduler state.
Taking #2516's conflict side would silently revert #2733. Its new positional args parameter also requires updating #2733/#2669 tests and callers.
Distinguish legacy warm start from a torn current checkpoint.
_optimizer_param_state_is_available maps both “no new-format files exist” and “only a subset exists” to the same weight-only warm start. The former is valid legacy compatibility; the latter is a damaged current checkpoint and should fail rather than silently discard optimizer, scheduler, and step state.
Enforce topology compatibility.
Adapter files become topology-keyed after #2669, while training/optimizer files remain keyed by global rank. Persist world/TP/PP/EP/DP topology metadata and refuse exact restore when it differs; otherwise the same world size with a different layout can pair correct adapter shards with another coordinate's optimizer state.
Coordinate each remaining mutation boundary.
Scheduler loading is currently unguarded after optimizer mutation. Likewise, optimizer child save failures should be checked after each child so all ranks stop in lockstep, rather than waiting until every child has run. The parameter-state payload also round-trips with weights_only=True, map_location="cpu"; use that safer load mode where possible.
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
Follow-up review on the unchanged head (55dc933d6c9070940f854a88f9111f018d98853c). These findings are new relative to the first review.
Blocker — rank-local checkpoint generations and iterations are never compared.
Each rank independently loads training_state_rank{rank}.pt, and the only consensus covers whether torch.load raised. There is no equality check for checkpoint generation, iteration, or scheduler position. I reproduced this with two valid files: rank 0 contained iteration 10 and rank 1 iteration 11; _load_training_state returned [10, 11] without an error.
This can arise from an interrupted overwrite because native, training-state, and optimizer-state files are published independently with no generation manifest. Persist one generation ID plus the expected shard/file set, write a completion marker only after all required state succeeds, and require all ranks to agree on generation and iteration before any model/optimizer mutation.
--no-load-optim can leave the LR/WD scheduler at step zero.
The no-load branch returns the saved nonzero iteration without restoring opt_param_scheduler. Later, model.py skips opt_param_scheduler.step(iteration * global_batch_size) whenever use_checkpoint_opt_param_scheduler is true and iteration is nonzero, assuming the checkpoint restored scheduler state. In this path that assumption is false, so training resumes at the old iteration with a fresh scheduler.
Return whether scheduler state was actually restored, and reposition it from iteration whenever it was not.
Security — no-load paths still perform unrestricted pickle deserialization.
training_state_rank*.pt is loaded with weights_only=False before checking optimizer is None or args.no_load_optim. A user-supplied adapter directory can therefore execute pickle payloads even when optimizer loading was explicitly disabled. The real distributed-Adam training-state payload produced by this PR successfully round-tripped with weights_only=True, map_location="cpu" in local validation.
Use safe deserialization for the supported state schema, or split iteration/scheduler metadata from any payload that genuinely requires trusted pickle loading.
Validate supported process-group and optimizer-child layouts up front.
The legacy distributed-optimizer parameter-state API requires Megatron DP Gloo subgroups. Miles passes args.enable_gloo_process_groups into optimizer construction, so --disable-gloo-process-groups currently fails only at the first LoRA save interval. Reject this incompatible combination during argument validation rather than failing after training has begun.
The active-child compatibility guard is also one-directional: checkpoints written with optimizer_children are checked, but a checkpoint written while all children were active uses the optimizer key and bypasses the guard when the current run has a stub child. Record active child indices unconditionally and compare them in both directions before calling load_state_dict.
|
Thanks — this review led to a substantial rewrite. The PR is now narrowed to one problem: Adopted — the DP-root preflight. Confirmed against Megatron: Adopted — Adopted — distinguish "legacy checkpoint" from "torn checkpoint". All shards present → restore. None present → warn and Adopted — collectives must not sit below a per-rank early return. Re-reviewing, Not adopted — guarding the scheduler load. Nothing collective runs after it, so a rank that skips it cannot strand anyone. Adding a consensus there is cost with no failure mode behind it. Not adopted — per-child consensus around the parameter-state save. Stub-ness is rank-dependent, so a global all-reduce per child would itself deadlock whenever the ranks disagree about how many non-stub children they have. The save loop records the first error and reaches one consensus at the end instead. Not adopted — Deferred (listed in the PR body). Note on stacking: the failure-consensus helpers here duplicate Validation: 4 focused tests exercising round-trip through the DP root, the legacy warm-start, the torn-shard rejection, and the pre-scatter shape rejection; |
ae2a445 to
8e1587c
Compare
ad58f15 to
f4ac9dc
Compare
DistributedOptimizer keeps its fp32 master weights and Adam moments out of state_dict() -- Megatron saves them separately via save_parameter_state -- so a LoRA checkpoint that stored only state_dict() resumed with neither. Save them per data-parallel root and restore them through the dp-zero scatter, coordinating rank-local I/O failures so no rank enters a collective its peers skipped, and checking the buffer sizes before the scatter rather than inside it. Checkpoints written before this warm-start instead of failing. Also refresh the masters in load_lora_adapter whenever adapter weights are written. They are snapshotted at optimizer construction, which happens before the load, so without this the first step() copies the construction-time values back over the restored adapter -- the same reason the Inkling path already calls reload_model_params. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
f4ac9dc to
f2af2cb
Compare
f972fa3 to
3706ef5
Compare
…atron chained APIs
# Conflicts: # miles/backends/megatron_utils/lora/utils.py
… master-refresh test
…nt_dir requires A rank whose adapter-shard write failed left save_parameter_state's data-parallel gather unmatched, so the save hung instead of raising.
…ep; gather first
get_parameter_state_dp_zero raises KeyError('exp_avg') on an optimizer that has
never stepped, so a save before the first step failed. With no parameter state,
resume keeps a fresh optimizer, which is what an unstepped one is.
|
Verified that LoRA checkpoints now save and restore the previously missing FP32 master weights and Adam moments, with stub optimizers correctly skipped. |
This branch has not been deployed
4 participants
Part of #2705.
Problem
LoRA checkpoints save and restore only
optimizer.state_dict(). Megatron'sDistributedOptimizerkeeps the FP32 master weights and Adam moments in its parameter state instead, so a resume trains from uninitialized moments, and loading the adapter leaves the masters at their pre-load values. With expert-only LoRA, the stub child in theChainedOptimizermakes the save itself raise.Fix
state_dict. A checkpoint without it (saved bymainor before the first step) resumes with a fresh optimizer; a partial one is rejected.state_dict/load_state_dict.Tests
tests/fast/backends/megatron_utils: 292 passed (286 onmain).mainthe save raises with a stub, and without one 18546/58884 params are NaN after one resumed step.iter_0000002/adapter: trains to the end. Onmainthe second resumed step has a NaN grad norm.Overlaps #2661 in
load_lora_adapter; the master refresh must stay ahead of its weight-only return.