Repository navigation
Conversation
Add --save-trainable-params-only, which filters the model section(s) of a torch_dist checkpoint down to parameters with requires_grad=True before save. This lets a PEFT/LoRA-style fine-tuning run (frozen base model, a small number of trainable adapter parameters) write a checkpoint sized to the adapter instead of duplicating the frozen base weights on every save. The optimizer section is left untouched, since it already only holds state for whichever parameters were passed to the optimizer. Loading such a checkpoint onto a fully-initialized base model requires --dist-ckpt-strictness log_unexpected (or another *_unexpected value): the on-disk checkpoint is missing the frozen keys the model will request, and this repurposes Megatron-Core's existing StrictHandling support for that rather than adding new load-side code. The new flag is restricted to --ckpt-format torch_dist, since only that format's ShardedTensor/ShardedTensorFactory leaves preserve each parameter's requires_grad through checkpointing; the other formats call state_dict(keep_vars=False), which detaches every tensor and would make every parameter look frozen. This is enforced both at argument-parse time and again at the point the filter actually runs, since args.ckpt_format can be reassigned after validate_args() (e.g. via --ckpt-convert-format). Fixes NVIDIA#7270 Signed-off-by: Matthew Romano <matthewromano5087@gmail.com>
Three real bugs found by actually running save/load/resume round trips on
GPU (2x A100, TP=1/2, PP=1/2, with and without SwiGLU), not caught by the
unit tests alone:
- The torch_dist-format guard was only checked in validate_args(), which
runs once at CLI-parse time. args.ckpt_format can be reassigned after
that (e.g. --ckpt-convert-format), so a converted-format save with
--save-trainable-params-only would silently detach every tensor and
drop the entire model section -- including the adapter weights being
saved, not just the frozen base model. Re-check the effective format at
the point the filter actually runs and fail loudly instead.
- TE `_extra_state` entries were kept unconditionally regardless of
their sibling weight's frozen/trainable status, since a ShardedObject
carries no requires_grad signal of its own. A frozen module's
extra_state surviving while its weight/bias got dropped left a
structurally inconsistent checkpoint (extra_state present, weight
missing at the same path) that failed to load. Drop `_extra_state`
unconditionally instead, matching Megatron-Bridge's proven PEFT
checkpoint filter.
- Loading an adapter-only checkpoint back onto a matching base model
crashed in apply_factory_merges whenever a frozen parameter was
ShardedTensorFactory-wrapped (e.g. SwiGLU's linear_fc1): the load path
only filtered the on-disk checkpoint, and relied on
--dist-ckpt-strictness log_unexpected to reconcile the *load request*
(built from the fresh, unfiltered model) against it. That reconciliation
doesn't correctly propagate through un-built ShardedTensorFactory
requests. Fix by filtering the load request's model section the same
way the save filtered what was written -- the same technique
Megatron-Bridge's own (production, SwiGLU-tested) PEFT checkpoint
loading uses. --dist-ckpt-strictness log_unexpected is still required
in addition, since the optimizer section is deliberately left
unfiltered and a separate load-time shape-validation list needs
non-strict handling for a frozen parameter's optimizer-section entry.
Verified with a real GPTModel (get_model/get_megatron_optimizer, plain
Adam since DistributedOptimizer doesn't support any frozen parameters at
all -- a pre-existing, unrelated limitation) across TP in {1,2}, PP in
{1,2}, and gated_linear_unit in {True,False}: adapter-only checkpoints
are correctly smaller than full checkpoints, frozen base weights survive
a non-strict load untouched, and the adapter parameters and iteration
count are correctly restored.
Also adds a unit test for the _extra_state fix.
Signed-off-by: Matthew Romano <matthewromano5087@gmail.com>
…ointing
tests/unit_tests/test_checkpointing.py::TestFilterStateDictToTrainableParams
only exercised the save-side filter function in isolation against synthetic
ShardedTensor objects; it never actually saved or loaded a real checkpoint, so
it couldn't have caught any of the three bugs fixed in the previous commit.
Add a real TP/PP/SwiGLU-parametrized test that builds an actual GPTModel with
a frozen base and a trainable adapter on every layer, saves it adapter-only,
loads it back onto a matching (corrupted) base model, and confirms: the
adapter-only checkpoint is smaller than the full one, the frozen base weights
are untouched by the load, the adapter weights are restored, and (via
optimizer.state_dict() comparison) the adapter's optimizer momentum state is
correctly restored -- i.e. resuming PEFT training actually works, not just
loading weights. Run at TP in {1,2}, PP in {1,2}, and gated_linear_unit in
{True,False} (6 cases), matching the real save/load/resume verification run
manually on 2x A100 for this PR.
Signed-off-by: Matthew Romano <matthewromano5087@gmail.com>
roma5087
marked this pull request as ready for review
September 16, 2026 20:14
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
This branch has not been deployed
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
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.
Fixes #7270
What this does
Adds
--save-trainable-params-only, which filters atorch_distcheckpoint'smodel section(s) down to parameters with
requires_grad=Trueat save time, andfilters the load request the same way at load time. For PEFT/LoRA-style
fine-tuning (frozen base model, a small number of trainable adapter parameters),
this keeps the saved checkpoint proportional to the adapter instead of
duplicating the frozen base weights on every save, and lets that checkpoint be
loaded back onto a matching base model (
--dist-ckpt-strictness log_unexpectedrequired at load time; see the new config field's docstring).
Verification
Beyond the unit tests below, this was verified end-to-end on a real GPTModel
(2x A100) across
--tensor-model-parallel-sizein {1,2},--pipeline-model-parallel-sizein {1,2}, and--swiglu/gated_linear_unitin {True,False}: adapter-only checkpoints are smaller than full checkpoints,
frozen base weights survive a non-strict load untouched, adapter weights and
optimizer momentum state (
exp_avg/exp_avg_sq) are correctly restored onresume, and iteration count is correctly restored.
It was also verified against Megatron-Bridge's real, unmodified
LoRA/PEFTimplementation (imported directly, not reimplemented): Bridge's LoRA correctly
applies to a Megatron-Core
GPTModel, and a checkpoint saved by this PR'ssave_checkpoint(save_trainable_params_only=True)loads correctly via Bridge'sreal
apply_peft_adapter_filter_to_state_dictsubstituted intoload_checkpoint()— for Bridge's actual production configuration(
DistributedOptimizer, real training before checkpointing; Bridge's owntest_finetune_lora.pyGPU CI test already exercises this combination andpasses).
Tests added
tests/unit_tests/test_checkpointing.py::TestFilterStateDictToTrainableParams— pure-function unit tests for the save-side filter (frozen/trainable leaves,
ShardedTensorFactory,_extra_state, multi-PP-stage sections).tests/unit_tests/dist_checkpointing/test_peft_checkpointing.py— realGPTModel save → load → resume round trip, parametrized over TP/PP/SwiGLU (6
cases), including optimizer momentum-state restoration.
Known limitations (not fixed in this PR)
--use-distributed-optimizerdoes not support any frozen parameters atall today (
DistributedOptimizer._build_optimizer_group_rangesraises aKeyErrorif a frozen param's grad-buffer bucket wasn't built the same waythe optimizer's param groups were) — a pre-existing Megatron-Core limitation,
unrelated to this change, that a real PEFT run must work around either by
freezing parameters before the model is DDP-wrapped (so the grad buffer and
optimizer param groups agree from construction) or by using the standard
wrap_with_ddp=Truetraining entrypoint path.a model with a
ShardedTensorFactory-wrapped trainable parameter (e.g. aLoRA adapter on a SwiGLU-gated
linear_fc1), can fail to load withapply_factory_merges: Different dict keys encountered.save_checkpointnever calls
init_state_fn(so a never-steppedtorch.optim.Adam'sstateis empty and nothing is written for the optimizer section), whileload_checkpointalways does (so a fresh optimizer's load request expectszero-seeded state for every trainable param) — for an ordinary
ShardedTensorthis asymmetry is tolerated by
--dist-ckpt-strictness log_unexpected, butapply_factory_merges's key-set equality check has no such escape hatch.This is a pre-existing
dist_checkpointing/optimizer gap independent of PEFTor this flag specifically (it would affect any GLU-gated model's iteration-0
checkpoint under the plain optimizer path), not something this PR introduces,
and does not affect a normal training run (which always takes at least one
step before its first checkpoint save).