Summary
When a LoRA experiment's checkpoint.load_path is a DCP-style URI without a .pt suffix (e.g. the checkpoint-DB / s3://...iter_000023000 default), model_loader.py classifies the checkpoint as "dcp" directly from the path string (cosmos_predict2/_src/predict2/utils/model_loader.py, line 205 on main):
checkpoint_format = "pt" if s3_checkpoint_dir.endswith(".pt") else "dcp"
But get_checkpoint_path() (line 219) resolves that URI through the checkpoint DB to a local consolidated .pt file. The LoRA key-mapping branch — which remaps plain checkpoint keys onto the PEFT-wrapped module names (base_model.model.*, base_layer.*) — is gated on checkpoint_format == "pt", so with a DCP-style load_path it is skipped. With use_lora=True the net's keys all carry PEFT prefixes, the subsequent non-strict load matches zero parameters, no exception is raised, and training proceeds from a randomly-initialized base.
The run looks healthy — loss decreases (the model is learning from scratch), checkpoints save normally — but the base model was never loaded. Inference from such a checkpoint produces noise, or output indistinguishable from the base model. This matches the symptoms in #176 ("LoRA fine-tuned model produces identical inference results to base model + checkpoint loading warnings").
Reproduction
- Run a LoRA post-training experiment whose
checkpoint.load_path is a checkpoint-DB URI without a .pt suffix, single node, use_lora=True.
- At startup,
debug.log shows load model in non-strict mode followed by _IncompatibleKeys(missing_keys=[...]) listing every base parameter (base_model.model.x_embedder.proj.1.weight, …) — and no "mapping checkpoint keys for LoRA" line.
- First-iteration training loss is ~3.0 (random base) instead of ~0.03 (loaded base).
- Generated video from the resulting checkpoint is noise.
Fix (one conditional)
Derive the format from the resolved file rather than the original URI, immediately after resolution (after line 219):
local_s3_ckpt_fp = get_checkpoint_path(cur_key_ckpt_full_path)
# The checkpoint DB can resolve a DCP-style URI to a consolidated .pt file.
# The LoRA key-mapping branch below is gated on checkpoint_format == "pt",
# so derive the format from the resolved file — otherwise a PEFT-wrapped
# model silently loads zero parameters and trains from random init.
if str(local_s3_ckpt_fp).endswith(".pt"):
checkpoint_format = "pt"
Verified on Cosmos-Predict2.5-2B LoRA post-training (rank 32, single H100): with the patch, startup logs Mapped 689 LoRA keys from checkpoint to model, first-iteration loss drops from ~3.04 to ~0.03, and generation produces coherent video. PR attached.
Two adjacent observations (happy to file separately)
- In LoRA runs the EMA weights are never synced from the loaded base (the EMA branch of the same load reports all-missing keys even after this fix), so EMA-derived consolidated checkpoints from LoRA runs are unusable — prefer the regular (
net) weights at conversion.
- A small diagnostic that diffs a converted checkpoint's frozen base weights against the official base catches this entire failure class in ~3 minutes on CPU; glad to contribute it if useful.
Summary
When a LoRA experiment's
checkpoint.load_pathis a DCP-style URI without a.ptsuffix (e.g. the checkpoint-DB /s3://...iter_000023000default),model_loader.pyclassifies the checkpoint as"dcp"directly from the path string (cosmos_predict2/_src/predict2/utils/model_loader.py, line 205 onmain):But
get_checkpoint_path()(line 219) resolves that URI through the checkpoint DB to a local consolidated.ptfile. The LoRA key-mapping branch — which remaps plain checkpoint keys onto the PEFT-wrapped module names (base_model.model.*,base_layer.*) — is gated oncheckpoint_format == "pt", so with a DCP-styleload_pathit is skipped. Withuse_lora=Truethe net's keys all carry PEFT prefixes, the subsequent non-strict load matches zero parameters, no exception is raised, and training proceeds from a randomly-initialized base.The run looks healthy — loss decreases (the model is learning from scratch), checkpoints save normally — but the base model was never loaded. Inference from such a checkpoint produces noise, or output indistinguishable from the base model. This matches the symptoms in #176 ("LoRA fine-tuned model produces identical inference results to base model + checkpoint loading warnings").
Reproduction
checkpoint.load_pathis a checkpoint-DB URI without a.ptsuffix, single node,use_lora=True.debug.logshowsload model in non-strict modefollowed by_IncompatibleKeys(missing_keys=[...])listing every base parameter (base_model.model.x_embedder.proj.1.weight, …) — and no "mapping checkpoint keys for LoRA" line.Fix (one conditional)
Derive the format from the resolved file rather than the original URI, immediately after resolution (after line 219):
Verified on Cosmos-Predict2.5-2B LoRA post-training (rank 32, single H100): with the patch, startup logs
Mapped 689 LoRA keys from checkpoint to model, first-iteration loss drops from ~3.04 to ~0.03, and generation produces coherent video. PR attached.Two adjacent observations (happy to file separately)
net) weights at conversion.