[2/N] Support lora checkpoint on partial save and multi-source load

[mori360]

Summary

Consolidate converter-checkpoint integration from 5 scattered dynamic attributes into a single ConverterCheckpointHooks dataclass, simplifying ModelWrapper and eliminating closure factories in LoRA.

Design: ConverterCheckpointHooks

  class ConverterCheckpointHooks:
      key_filter: Callable[[str], bool] | None = None
      save_last_fn: Callable | None = None
      load_additional_fn: Callable | None = None
      finalize_fn: Callable | None = None

• key_filter — Distinguishes converter-added keys (e.g. LoRA's .lora_a., .lora_b.) from base model keys. Used by
  ModelWrapper.state_dict() to save only adapter keys in periodic checkpoints (small, fast saves), and by base_state_dict() to exclude them when loading base weights from HF. Without this, every periodic checkpoint would save the full model.
• save_last_fn — Converter-specific save format for the final checkpoint. LoRA uses this to write PEFT-compatible adapter_model.safetensors + adapter_config.json so the result can be loaded directly by HuggingFace PEFT/vLLM. Only called when last_save_in_hf=True. Without this, the final save would use DCP format which external tools can't read.
• load_additional_fn — Loads converter-specific checkpoint formats during resume. LoRA uses this to load PEFT safetensors and remap keys back to torchtitan naming. Called for secondary checkpoint paths (the adapter checkpoint, separate from the base model HF checkpoint). Without this, resuming from a PEFT save would fail.
• finalize_fn — End-of-training model transformation before the last save. Runs ModelConvertersContainer.finalize() in reverse order: LoRA merge (fold adapters into base weights) then QAT CONVERT (replace fake-quantized modules with real quantized ones). This is set by ModelConvertersContainer.convert(), not by individual converters. Without this, the final checkpoint would contain training-time artifacts (LoRA adapters, fake quantization) instead of a clean deployable model.

Changes

• model_converter.py: Added ConverterCheckpointHooks dataclass. ModelConvertersContainer.convert() merges into existing hooks.
• lora.py: Replaced closure factories (_make_peft_save_fn, _make_peft_load_fn, _make_merge_fn) with regular methods (_save_peft, _load_peft). convert() builds hooks directly. finalize() cleans up hooks after merge.
• checkpoint.py: ModelWrapper simplified from 7 methods to 5. Added get_hooks() as single access point. Removed scattered getattr(part, "converter*") reads.
• qat.py: Extracted apply_qat_prepare() shared by QATConverter.convert() and LoRAConverter._apply_adapter_qat().

Test plan

• pytest tests/unit_tests/test_checkpoint.py -x — new TestModelWrapperConverterKeys tests (strict vs
  partial planner)
• pytest tests/unit_tests/test_model_converter.py -x — new test_lora_key_remap_roundtrip
• torchrun --nproc_per_node=4 -m torchtitan.train --module llama3 --config llama3_debugmodel_lora —
  verify LoRA training runs end-to-end
• torchrun --nproc_per_node=4 -m torchtitan.train --module llama3 --config llama3_8b_lora — verify 8B
  LoRA config with PEFT save

Stack from ghstack (oldest at bottom):

• [5/N] Support lora moe #2569
• [4/N] Add QAT (quantization-aware training) model converter #2488
• -> [2/N] Support lora checkpoint on partial save and multi-source load #2485
• [1/N] Add lora adapters for finetune #2484

[tianyu-l]

The new truth table is way too complicated. I have doubts in both

• whether we need to support all various combinations
• if so, whether this is the best flow to support them

Let's start with the following questions:
When user train lora, do they ever need to store anything other than the adapters?

For fault tolerance or changing hyperparameters, users might want to save checkpoint and resume training from previous checkpoint.

• With small scale, loading main weights in HF format, and loading lora weight in HF / DCP is good enough.
• With large scale, for main weight DCP may have advantage over HF, but at large scale why would user use LoRA in the first place?

For inference, I think vllm supports loading LoRA adapters, so again we only need to save lora weights. Please check what format vllm accepts.

So, is it true that we only need to differentiate two cases:

• using LoRA to train, but it's the first time we load an HF checkpoint
• continued LoRA training, where we need to load HF checkpoint and LoRA weights

[fegin]

With small scale, loading main weights in HF format, and loading lora weight in HF / DCP is good enough.

Do we have the number for this? iirc, the conversion is not cheap though it may be acceptable on a single machine.

And this doesn't simplify the logic as we will need to save trainer/optimizer states for fault tolerance purpose. The original design is coupled that with the model weights, which are frozen now.

If we are okay with always loading with HF format, even for the fault tolerance ones, we should treat the HF folder as the additional folder. lora adapter + trainer + optimizer are in the original checkpoint folder. cc., @mori360