Merge LoRA adapters safely in automatic Hugging Face exports

.github/workflows/gpu-tests.yml

@@ -523,6 +523,12 @@ jobs:
           uv run --all-extras --group=cu128-train python -m pytest -v -s \
             "$test_file::$test_name" -o addopts= -m L0
 
+      - name: Distributed unit tests - HF LoRA export (torchrun, 2 ranks)
+        run: |
+          export LD_LIBRARY_PATH=
+          uv run --all-extras --group=cu128-train torchrun --nproc_per_node=2 -m pytest -v \
+            cosmos_framework/callbacks/hf_export_test.py -o addopts= -m L0 -k fsdp
+
       # The cfgp_ar / context_parallel tests call dist.init_process_group and
       # skip under plain pytest (world_size 1); they must be launched with
       # torchrun. They hardcode their world size via ParallelDims, so each file

cosmos_framework/callbacks/hf_export.py

@@ -32,6 +32,7 @@
 from cosmos_framework.utils import log
 from cosmos_framework.utils.callback import Callback
 from cosmos_framework.utils.distributed import is_rank0
+from cosmos_framework.utils.generator.lora import LoraInjectedLinear
 
 try:
     from safetensors.torch import save_file as _safetensors_save_file
@@ -53,6 +54,56 @@
 }
 
 
+def _hf_parameter_name(name: str) -> str:
+    """Remove complete compile/checkpoint wrapper segments, including at the end."""
+    return ".".join(part for part in name.split(".") if part not in {"_orig_mod", "_checkpoint_wrapped_module"})
+
+
+def _lora_merge_plan(network: torch.nn.Module) -> tuple[dict[str, LoraInjectedLinear], set[str]]:
+    # Inspect every path before named_parameters() deduplicates shared weights.
+    # A plain path and an adapted path (or two distinct adapters) cannot share a
+    # single dense export weight: their effective weights may differ. Repeated
+    # aliases of the same LoRA module are safe, as are ordinary non-LoRA ties.
+    lora_owners = {
+        f"{name}.weight" if name else "weight": module
+        for name, module in network.named_modules(remove_duplicate=False)
+        if isinstance(module, LoraInjectedLinear)
+    }
+    aliases: dict[int, list[str]] = {id(module.weight): [] for module in lora_owners.values()}
+    if aliases:
+        for name, parameter in network.named_parameters(remove_duplicate=False):
+            if id(parameter) in aliases:
+                aliases[id(parameter)].append(name)
+        ambiguous = [
+            sorted(names)
+            for names in aliases.values()
+            if any(name not in lora_owners for name in names) or len({id(lora_owners[name]) for name in names}) > 1
+        ]
+        if ambiguous:
+            # This topology-only preflight runs on every rank before any tensor
+            # all-gather, so rejection does not leave peers waiting in a gather.
+            raise ValueError(f"Cannot export ambiguous shared LoRA base weights: {sorted(ambiguous)}")
+
+    merges: dict[str, LoraInjectedLinear] = {}
+    adapter_keys: set[str] = set()
+    for name, module in network.named_modules():
+        if isinstance(module, LoraInjectedLinear):
+            prefix = _hf_parameter_name(name)
+            prefix = f"{prefix}." if prefix else ""
+            weight_name = f"{prefix}weight"
+            if weight_name in merges:
+                raise ValueError(f"Multiple LoRA modules map to HF parameter {weight_name!r}")
+            merges[weight_name] = module
+            adapter_keys.update((f"{prefix}lora_A.weight", f"{prefix}lora_B.weight"))
+    return merges, adapter_keys
+
+
+def _gather_full(parameter: torch.Tensor) -> torch.Tensor:
+    if isinstance(parameter, torch.distributed.tensor.DTensor):
+        parameter = parameter.full_tensor()
+    return parameter.detach()
+
+
 def _upload_folder_to_s3(local_folder: str, bucket: str, s3_prefix: str, credential_path: str) -> None:
     """Upload every file under *local_folder* to ``s3://{bucket}/{s3_prefix}/...``.
 
@@ -211,6 +262,10 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di
         ``cpu_chunks`` and ``manifest``; other ranks return empty structures but still
         participate in the distributed all-gathers.
 
+        LoRA updates are merged into detached export tensors. Adapter parameters
+        are gathered immediately after their base weight on every rank and are
+        omitted from the dense HF checkpoint. Live training weights stay intact.
+
         Returns:
             cpu_chunks:  List of ``{weight_name: cpu_tensor}`` dicts, one per shard file.
             manifest:    Mapping of ``weight_name → shard_filename``.
@@ -223,7 +278,13 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di
         total_size: int = 0
         file_idx: int = 0
 
-        for name, param in model.model.model.named_parameters():
+        network = model.model.model
+        merges, adapter_keys = _lora_merge_plan(network)
+        merged_names: set[str] = set()
+        exported_names: set[str] = set()
+        duplicate_names: set[str] = set()
+
+        for name, param in network.named_parameters():
             # Phase 2+: HFModel initialises _model via AutoModelForImageTextToText /
             # AutoModelForCausalLM, so parameter names are HF-native and match the
             # safetensors checkpoint keys loaded by _load_vlm_weights().
@@ -239,14 +300,23 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di
             # torch.compile and gradient-checkpointing wrappers inject prefixes into
             # named_parameters() output.  Strip them so exported keys are HF-native,
             # matching what HFModel._load_vlm_weights() does for the in-memory state dict.
-            name = name.replace("_orig_mod.", "").replace("_checkpoint_wrapped_module.", "")
+            name = _hf_parameter_name(name)
+            if name in adapter_keys:
+                continue  # Already gathered at this adapter's base weight.
 
             # Gather across FSDP / TP / CP ranks (collective — all ranks must call).
-            if isinstance(param, torch.distributed.tensor.DTensor):
-                param = param.full_tensor()
-            param = param.detach()
-            if self._export_dtype is not None:
+            param = _gather_full(param)
+            if name in merges:
+                module = merges[name]
+                lora_a = _gather_full(module.lora_A.weight)
+                lora_b = _gather_full(module.lora_B.weight)
+                param = module.merged_weight(param, lora_a, lora_b, dtype=self._export_dtype)
+                merged_names.add(name)
+            elif self._export_dtype is not None:
                 param = param.to(dtype=self._export_dtype)
+            if name in exported_names:
+                duplicate_names.add(name)
+            exported_names.add(name)
 
             tensor_bytes = param.element_size() * param.numel()
 
@@ -262,11 +332,22 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di
 
             shard_name = f"{file_idx:05d}.safetensors"
             if is_rank0():
-                current_chunk[name] = param.cpu()
+                # CPU parameters also need independent storage: the background
+                # writer must not observe subsequent optimizer updates.
+                current_chunk[name] = param.to(device="cpu", copy=True)
                 manifest[name] = shard_name
                 total_size += tensor_bytes
             current_chunk_bytes += tensor_bytes
 
+        # After all collectives, ensure every planned merge was exported and
+        # normalized names stayed unambiguous. Keep this check on every rank.
+        missing_merges = merges.keys() - merged_names
+        if missing_merges or duplicate_names:
+            raise ValueError(
+                "Cannot export an unambiguous dense HF checkpoint: "
+                f"unmerged LoRA weights={sorted(missing_merges)}, duplicate keys={sorted(duplicate_names)}"
+            )
+
         # Flush the final (possibly partial) chunk.
         if current_chunk_bytes > 0 and is_rank0() and current_chunk:
             cpu_chunks.append(current_chunk)