feat(checkpoint): support allow_partial_load_on_model

[chango9543]

Summary

Resolves #4010.

Adds a checkpoint.allow_partial_load_on_model config field that enables loose checkpoint loading for the model portion of the state dict only. When True, model keys present in the in-memory state dict but absent from the checkpoint are skipped (left at their current / random init) instead of raising a runtime error. Non-model state (optimizer / lr_scheduler / dataloader / train_state) is still loaded strictly — a missing key there still raises, matching today's behavior. Skipped model keys are logged at WARNING level after the load completes.

This is the landing of the allow_partial_load_on_model compromise proposed by @ivy-zhou in #4010, addressing @sdmyzlp's review points on naming, scope, and transparency.

Motivation

When resuming or fine-tuning a model whose architecture has changed since the checkpoint was written (e.g. new transformer blocks, a new head, extra MoE experts, adapters), users want to load the pretrained backbone and leave the new parts at random init. Today CheckpointManager.dcp_load calls torch.distributed.checkpoint.load(...) without a custom planner, which uses DefaultLoadPlanner() with allow_partial_load=False (strict), so any key present in the model but absent from the checkpoint raises. This PR exposes the existing DCP knob (DefaultLoadPlanner(allow_partial_load=True)) but scopes it to model keys so optimizer/lr_scheduler state is never silently dropped.

Design

• Config-driven (no method signature changes): a new checkpoint.allow_partial_load_on_model: bool = False field, read in __init__ alongside exclude_from_loading and the other loading-policy fields.
• In dcp_load, three branches sharing a common skipped-keys logging epilogue:
  • HF path (from_hf=True): uses DefaultLoadPlanner(allow_partial_load=self.allow_partial_load_on_model) directly. This is safe because initial_load_in_hf forces initial_load_model_only=True (enforced in Config.__post_init__), so the state_dict only contains model keys — there is no non-model state to protect.
  • DCP path + flag True: splits the flattened state dict into model_sd (flat model FQNs, loaded with DefaultLoadPlanner(allow_partial_load=True)) and non_model_sd (top-level state groups, loaded with the default strict DefaultLoadPlanner()).
  • DCP path + flag False (default): original strict behavior, single dcp.load.
• After the load, skipped keys are computed once via set(model_planner.state_dict) - set(model_planner.metadata.state_dict_metadata) and logged via logger.warning. This epilogue applies to both the HF and DCP-loose paths; in the strict-default path model_planner.metadata is None (the planner isn't used) so the epilogue is a no-op.

Why config-driven instead of a dcp_load(..., allow_partial_load=...) parameter

Matches the existing pattern — exclude_from_loading, initial_load_model_only, load_step are all Config fields consumed as self.* inside the manager, not threaded through method signatures. A loading policy is a deployment-level decision, which fits a Config field better than a per-call param.

How @sdmyzlp's review points are addressed

• Naming → allow_partial_load (matches DefaultLoadPlanner's parameter name; the skipped keys live in the model state dict, not the checkpoint, so the direction is correct). The config field is named allow_partial_load_on_model to reflect the scoped-to-model semantics.
• Scope → loose loading is scoped to the model key only; non-model state stays strict (raises on missing keys), so optimizer/lr_scheduler state is never silently dropped.
• Transparency → skipped model keys are logged via logger.warning after the load, using set(planner.state_dict) - set(planner.metadata.state_dict_metadata) (the flattened namespace DCP actually planned against).

Proof of Value

Loss

This change does not affect computation results. It only modifies the checkpoint load path (dcp_load); training forward/backward is untouched. With allow_partial_load_on_model=False (default), load behavior is identical to today, so loss is unchanged. With the flag True, only which checkpoint keys get copied into the model changes (newly-added module keys stay at random init instead of raising), which is the user's explicit intent — the training computation graph itself is unmodified.

Performance

No throughput/memory impact on training. The only overhead is a one-time state-dict split and an extra dcp.load call (two calls instead of one) at load time, which is negligible relative to checkpoint IO.

Backward compatibility

• Default is False; existing behavior unchanged.
• No on-disk checkpoint format changes.
• No method signature changes — only a new Config field read in __init__ and a branch in dcp_load.

Testing

Added unit tests in tests/unit_tests/test_checkpoint.py (TestCheckpointManager):

Test	Verifies
test_allow_partial_load_on_model_skips_new_model_keys	Model keys absent from checkpoint are skipped (left at init), existing keys loaded
test_allow_partial_load_on_model_strict_for_missing_optimizer	Non-model keys absent from checkpoint still raise (strict gate holds)

All tests mock dcp.load and torch.distributed.get_rank following the existing patterns in test_checkpoint.py, so they run without a GPU or a real distributed process group.

Limitations / Follow-ups

• The HF path relies on the initial_load_in_hf ⟹ initial_load_model_only invariant (enforced in Config.__post_init__); if that invariant is ever relaxed, the HF path would need to split model/non-model like the DCP path does.
• torch_checkpointing migration: @ivy-zhou noted torchtitan is moving to torch_checkpointing; this feature can migrate cleanly into the new system when that lands.

Checklist

• Forked repo and branch from main
• Added tests for the new behavior
• Updated the Config docstring (API doc)

[meta-cla]

Hi @chango9543!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[pytorch-bot]

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

• ciflow/8gpu

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[ivy-zhou]

I think it might be nicer to create like a ModelPartialLoadPlanner, perhaps, so that we can have missing keys fail in planning phase before tensor load? Thoughts cc: @fegin ?

[ivy-zhou]

draft from LLM:

class ModelPartialLoadPlanner(DefaultLoadPlanner):
      """Allow missing model keys while keeping auxiliary state strict."""

      def __init__(self, model_state_keys: set[str]) -> None:
          super().__init__(allow_partial_load=True)
          self._model_state_keys = model_state_keys
          self.missing_model_fqns: tuple[str, ...] = ()

      def create_local_plan(self) -> LoadPlan:
          # Must run first: DCP may switch to pre-2.4 flattening here.
          plan = super().create_local_plan()
          assert self.metadata is not None

          missing = set(self.state_dict) - set(
              self.metadata.state_dict_metadata
          )

          # mappings[fqn][0] identifies the original top-level state-dict key.
          missing_model = {
              fqn
              for fqn in missing
              if self.mappings[fqn][0] in self._model_state_keys
          }
          missing_non_model = missing - missing_model

          if missing_non_model:
              raise RuntimeError(
                  "Missing non-model keys in checkpoint: "
                  f"{sorted(missing_non_model)}"
              )

          # DCP gathers local plans before calling create_global_plan.
          return replace(
              plan,
              planner_data=tuple(sorted(missing_model)),
          )

      def create_global_plan(self, plans: list[LoadPlan]) -> list[LoadPlan]:
          missing: set[str] = set()

          for plan in plans:
              assert isinstance(plan.planner_data, tuple), (
                  "StorageReader did not preserve planner_data"
              )
              missing.update(plan.planner_data)

          self.missing_model_fqns = tuple(sorted(missing))
          return super().create_global_plan(plans)