fix(lora): persist distributed optimizer parameter state - #2516
Conversation
6997c74 to
de6fb63
Compare
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Arist12
changed the title
ce1b28d to
d651c3c
Compare
d651c3c to
12f7112
Compare
Registering a new 900s 4-GPU E2E is a CI-budget decision that belongs with the reviewer, not with an optimizer-state fix. The save/load helpers are already covered by the fast tests in this PR, and exact resume was validated by hand on 8x MI350X.
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.
…cope Drop the optimizer-metadata, resume-flag and adapter-read changes, which fix different bugs, and select DistributedOptimizer children by type so the load path keeps its signature. Read the parameter state with weights_only and check it against the grad buffers on the DP root, before the scatter that would otherwise hang the group on a mismatch.
|
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; |
2 participants
Part of #2705.
Problem
DistributedOptimizershards the FP32 master weights and the Adam moments across the data-parallel group and deliberately keeps them out ofstate_dict(); they are saved separately throughsave_parameter_state/load_parameter_state_from_dp_zero. The LoRA checkpoint path only wroteoptimizer.state_dict(), so a resumed run restored the param-group metadata but rebuilt the master weights from the base model — the loaded adapter is then wiped by the first optimizer step, and the moments restart at zero.Change
save_lora_checkpointalso writesoptimizer_param_state_rank{rank}_optimizer{i}.ptfor every non-stubDistributedOptimizerchild, andload_lora_adapterrestores it. Both Megatron calls are collective (a gather / a scatter over the DP group), so the surrounding logic keeps every rank on the same branch:torch.saveresult is agreed on before the parameter-state gather, so a rank that fails to write cannot leave its peers in the gather;numel_unpaddedis compared against the grad buffers before any scatter.load_parameter_state_from_dp_zeroasserts the same thing between two scatters, where a DP-root failure hangs the rest of the group instead of raising.Checkpoints written before this change have no parameter-state files at all: they log a warning and
reload_model_params(), which at least seeds the masters from the adapter that was just loaded. A partially present set is an error rather than a silent warm start.The parameter state is read with
weights_only=True; the training state still needs full unpickling for its param-group metadata.Validation
test_lora_optimizer_checkpoint.py,test_lora_utils.py); fulltests/fast/backends/megatron_utilsgreen apart from twoTestUpdateWeightsZeroChunksfailures that reproduce onmain.Notes for reviewers
isinstance(child, DistributedOptimizer)rather thanargs.use_distributed_optimizer, which keepsload_lora_adapter's signature unchanged. OnlyDistributedOptimizer(andChainedOptimizer, which is never a child) definessave_parameter_state.DistributedOptimizer.load_parameter_state, which early-returns onis_stub_optimizer.--no-save-optim/--no-load-optim/--finetuneare not honoured by the LoRA checkpoint path at all (pre-existing, and fixing it changes the behaviour oftraining_state_rank*.pt);optimizer.state_dict()still raises when a chain has stub children; and_all_ranks_true/_raise_if_any_rank_failedduplicate helpers added by fix(lora): deduplicate native adapter checkpoints #2669 inlora_checkpoint.py, so whichever lands second should reuse that module.