Skip to content

LoRA post-training silently trains from random base weights when load_path is a checkpoint-DB URI (root cause of #176) #160

Description

@Rouhi-Amirreza

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

  1. Run a LoRA post-training experiment whose checkpoint.load_path is a checkpoint-DB URI without a .pt suffix, single node, use_lora=True.
  2. 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.
  3. First-iteration training loss is ~3.0 (random base) instead of ~0.03 (loaded base).
  4. 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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions