feat(vlm): LoRA SFT support for VLM + VideoPhy-2 LoRA recipes

.github/workflows/gpu-tests.yml

@@ -337,6 +337,16 @@ jobs:
           uv run --all-extras --group=cu128-train torchrun --nproc_per_node=4 -m pytest -v \
             cosmos_framework/model/generator/mot/context_parallel_test.py -o addopts=
 
+      # HFExportCallback folds LoRA adapters into the base weights and gathers
+      # lora_A / lora_B from inside the base weight's loop iteration, so the
+      # ordering is only exercised once the adapters are DTensors. World size is
+      # fixed at 2; the test skips itself under any other.
+      - name: Distributed unit tests - hf_export LoRA merge (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_fsdp_test.py -o addopts=
+
       # Clear everything the suite writes into the working tree (all gitignored
       # scratch): pytest tmp dirs (DCP checkpoint, logs), the script-test
       # `outputs/` dir, any `examples/checkpoints`, and the `schemas/` dir from

cosmos_framework/callbacks/hf_export.py

@@ -12,6 +12,11 @@
 - Worker exceptions are stored in ``_worker_exception`` and re-raised on the next
   checkpoint or at train end, so failures are never silently swallowed.
 - Controlled entirely via ``config.checkpoint.hf_export`` (HFExportConfig).
+- LoRA runs export a MERGED checkpoint: each ``LoraInjectedLinear`` contributes a
+  single ``<path>.weight`` equal to ``W + (alpha / r) * B @ A``, and no ``lora_*``
+  keys are written. The export is therefore a plain HF checkpoint that
+  ``from_pretrained`` (and ``eval_videophy2``) loads with the adapter's effect
+  already in the weights, exactly like a full fine-tune's export.
 
 Phase 2+ note
 -------------
@@ -205,24 +210,89 @@ def on_train_end(self, model: Any, iteration: int = 0) -> None:
     # Internal helpers
     # ------------------------------------------------------------------
 
+    # Path segments torch.compile and gradient checkpointing insert into the
+    # module tree. They are not part of the HF-native name.
+    _WRAPPER_SEGMENTS: frozenset[str] = frozenset({"_orig_mod", "_checkpoint_wrapped_module"})
+
+    @classmethod
+    def _strip_wrapper_prefixes(cls, name: str) -> str:
+        """Drop wrapper segments from a parameter or module path.
+
+        Dropping whole dot-separated segments rather than substrings matters for
+        module paths: a wrapped module's own path *ends* with the wrapper segment
+        (``layer._checkpoint_wrapped_module``) and has no trailing dot to match
+        on, so substring removal would leave it alone and
+        :meth:`_lora_merge_plan` would key its adapters off a name that no
+        stripped parameter ever produces — a silently unmerged export.
+        """
+        return ".".join(seg for seg in name.split(".") if seg not in cls._WRAPPER_SEGMENTS)
+
+    @staticmethod
+    def _lora_merge_plan(root: torch.nn.Module) -> tuple[dict[str, Any], set[str]]:
+        """Locate every LoRA-adapted linear and the adapter keys it owns.
+
+        Returns ``(merge_targets, adapter_keys)`` where ``merge_targets`` maps a
+        base-weight parameter name to the ``LoraInjectedLinear`` holding it, and
+        ``adapter_keys`` is the set of ``lora_A`` / ``lora_B`` parameter names
+        that must NOT be written to the export. Both use post-strip names so
+        they match what :meth:`_gather_weights` computes.
+
+        Empty on a full fine-tune, which is what keeps that path untouched.
+        """
+        # Deferred: cosmos_framework.utils.generator.lora is only needed when a
+        # LoRA run reaches export, and hf_export is imported from config land.
+        from cosmos_framework.utils.generator.lora import LoraInjectedLinear
+
+        merge_targets: dict[str, Any] = {}
+        adapter_keys: set[str] = set()
+        for module_name, module in root.named_modules():
+            if not isinstance(module, LoraInjectedLinear):
+                continue
+            path = HFExportCallback._strip_wrapper_prefixes(module_name)
+            # An adapted linear at the tree root (or under nothing but wrappers)
+            # strips to "", and its parameters are plain "weight" / "lora_A.weight".
+            prefix = f"{path}." if path else ""
+            merge_targets[f"{prefix}weight"] = module
+            adapter_keys.add(f"{prefix}lora_A.weight")
+            adapter_keys.add(f"{prefix}lora_B.weight")
+        return merge_targets, adapter_keys
+
+    @staticmethod
+    def _gather_full(param: torch.Tensor) -> torch.Tensor:
+        """All-gather a sharded parameter. Collective — every rank must call it."""
+        if isinstance(param, torch.distributed.tensor.DTensor):
+            param = param.full_tensor()
+        return param.detach()
+
     def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], dict[str, str], int]:
         """Iterate model parameters, all-gather DTensor shards, and build CPU chunks.
 
         Must be called on **all ranks**. Only rank 0 populates the returned
         ``cpu_chunks`` and ``manifest``; other ranks return empty structures but still
         participate in the distributed all-gathers.
 
+        LoRA adapters are merged into their base weights here, so the export is a
+        plain HF checkpoint either way — see :meth:`_lora_merge_plan`.
+
         Returns:
             cpu_chunks:  List of ``{weight_name: cpu_tensor}`` dicts, one per shard file.
             manifest:    Mapping of ``weight_name → shard_filename``.
             total_size:  Total byte count of all exported tensors (for the index JSON).
         """
+        merge_targets, adapter_keys = self._lora_merge_plan(model.model.model)
+        if merge_targets:
+            log.info(
+                f"[HFExportCallback] Merging {len(merge_targets)} LoRA adapter(s) into their "
+                "base weights; the export carries no lora_* keys."
+            )
+
         cpu_chunks: list[dict[str, torch.Tensor]] = []
         manifest: dict[str, str] = {}
         current_chunk: dict[str, torch.Tensor] = {}
         current_chunk_bytes: int = 0
         total_size: int = 0
         file_idx: int = 0
+        merged: set[str] = set()
 
         for name, param in model.model.model.named_parameters():
             # Phase 2+: HFModel initialises _model via AutoModelForImageTextToText /
@@ -240,12 +310,32 @@ 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 = self._strip_wrapper_prefixes(name)
+
+            # Adapter tensors are folded into their base weight below, so they must
+            # not also be written out: no HF architecture declares lora_* keys, and
+            # from_pretrained() drops unexpected ones with a warning — the export
+            # would look complete while actually being the untuned base model.
+            if name in adapter_keys:
+                continue
 
             # 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()
+            param = self._gather_full(param)
+
+            lora_module = merge_targets.get(name)
+            if lora_module is not None:
+                # lora_A / lora_B are gathered here instead of at their own
+                # named_parameters() entries. Every rank walks the same module
+                # tree in the same order, which is all the all-gather requires.
+                param = lora_module.merged_weight(
+                    param,
+                    self._gather_full(lora_module.lora_A.weight),
+                    self._gather_full(lora_module.lora_B.weight),
+                )
+                merged.add(name)
+
+            # Cast after the merge: merged_weight accumulates in float32, and
+            # casting first would round the delta away before it is added.
             if self._export_dtype is not None:
                 param = param.to(dtype=self._export_dtype)
 
@@ -272,6 +362,37 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di
         if current_chunk_bytes > 0 and is_rank0() and current_chunk:
             cpu_chunks.append(current_chunk)
 
+        # Every adapter the plan found must actually have been folded in. The plan
+        # keys off named_modules() paths and the loop off named_parameters() paths;
+        # if a wrapper this code does not know about ever desynchronizes the two,
+        # the merge silently no-ops and the export is the untuned base model.
+        #
+        # `merged` is tracked on every rank (the loop is rank-independent), so this
+        # aborts everywhere at the same point rather than on rank 0 alone. Both
+        # checks sit after the last collective, so raising cannot strand a peer
+        # mid-all-gather.
+        if len(merged) != len(merge_targets):
+            missed = sorted(set(merge_targets) - merged)
+            raise RuntimeError(
+                f"[HFExportCallback] LoRA merge incomplete: {len(merged)} of "
+                f"{len(merge_targets)} adapters folded in. Unmerged base weights: "
+                f"{missed[:8]}{' ...' if len(missed) > 8 else ''}. The module paths from "
+                "named_modules() no longer line up with the parameter paths from "
+                "named_parameters() — check _WRAPPER_SEGMENTS for a wrapper this code "
+                "does not strip."
+            )
+        # The invariant the export must actually satisfy, asserted directly.
+        # manifest is rank-0-only, so this is a rank-0 check; the count above is
+        # what catches the desync case on every rank.
+        leaked = sorted(k for k in manifest if "lora_" in k)
+        if leaked:
+            raise RuntimeError(
+                f"[HFExportCallback] Adapter tensors leaked into the export: {leaked[:8]}"
+                f"{' ...' if len(leaked) > 8 else ''}. An HF checkpoint must carry merged "
+                "weights only; from_pretrained() would drop these and hand back the "
+                "untuned base model."
+            )
+
         return cpu_chunks, manifest, total_size
 
     def _save_and_upload(

cosmos_framework/callbacks/hf_export_fsdp_test.py

@@ -0,0 +1,162 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Distributed counterpart to ``hf_export_test.py`` — LoRA merging under FSDP2.
+
+``hf_export_test.py`` runs single-process on CPU with plain ``nn.Linear``, so it
+cannot cover the part that actually carries risk: ``lora_A`` / ``lora_B`` are
+DTensors sharded across ranks, and ``_gather_weights`` all-gathers them from
+inside the *base weight's* loop iteration rather than at their own
+``named_parameters()`` entries. That reordering is safe only because every rank
+walks the module tree identically — a property worth asserting rather than
+arguing.
+
+World size must be 2. Launch with::
+
+    torchrun --nproc_per_node=2 -m pytest cosmos_framework/callbacks/hf_export_fsdp_test.py
+
+Under plain pytest (no ``RANK``) every test skips, matching ``cfgp_ar_test`` and
+``context_parallel_test``.
+"""
+
+import os
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import checkpoint_wrapper
+from torch.distributed.fsdp import fully_shard
+
+from cosmos_framework.callbacks.hf_export import HFExportCallback
+from cosmos_framework.utils.generator.lora import LoraInjectedLinear
+
+_WORLD_SIZE = 2
+
+
+def setup_distributed_environment() -> tuple[int, int]:
+    if "RANK" not in os.environ:
+        pytest.skip("requires distributed environment (run with: torchrun --nproc_per_node=2)")
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl", init_method="env://")
+    rank, world_size = dist.get_rank(), dist.get_world_size()
+    if world_size != _WORLD_SIZE:
+        pytest.skip(f"requires world_size={_WORLD_SIZE}, got {world_size}")
+    torch.cuda.set_device(rank)
+    return rank, world_size
+
+
+def _lora_linear(in_f: int, out_f: int, rank: int, alpha: int, *, bias: bool = False) -> LoraInjectedLinear:
+    """Materialized adapter — ``__init__`` puts lora_A / lora_B on the meta device."""
+    base = nn.Linear(in_f, out_f, bias=bias)
+    module = LoraInjectedLinear(base, rank, alpha)
+    module.lora_A = nn.Linear(in_f, rank, bias=False)
+    module.lora_B = nn.Linear(rank, out_f, bias=False)
+    return module
+
+
+class _Block(nn.Module):
+    """Four adapted projections plus an unadapted MLP."""
+
+    _ADAPTED = ("q_proj", "k_proj", "v_proj", "o_proj")
+
+    def __init__(self, dim: int = 64, rank: int = 8, alpha: int = 16) -> None:
+        super().__init__()
+        for name in self._ADAPTED:
+            setattr(self, name, _lora_linear(dim, dim, rank, alpha, bias=(name == "o_proj")))
+        self.mlp = nn.Linear(dim, dim * 2)
+
+
+def _build_sharded_block() -> tuple[nn.Module, dict[str, torch.Tensor]]:
+    """Return an FSDP2-sharded block and the merged weights computed BEFORE sharding.
+
+    The reference is the whole point: it is what a single-process export of the
+    same model would produce, so comparing against it catches any way the
+    sharded path could diverge.
+    """
+    torch.manual_seed(1234)  # identical init on every rank
+    model = _Block().cuda()
+    for name in _Block._ADAPTED:
+        module = getattr(model, name)
+        # lora_B ships zero-initialized; a zero adapter makes the merge a no-op
+        # and would let a broken merge pass.
+        nn.init.normal_(module.lora_A.weight, std=0.02)
+        nn.init.normal_(module.lora_B.weight, std=0.02)
+
+    reference: dict[str, torch.Tensor] = {}
+    for name, module in model.named_modules():
+        if isinstance(module, LoraInjectedLinear):
+            reference[f"{name}.weight"] = module.merged_weight(
+                module.weight.detach(), module.lora_A.weight.detach(), module.lora_B.weight.detach()
+            ).clone()
+    for name, param in model.named_parameters():
+        # Adapters are folded into the base weight, so a correct export does not
+        # carry them and neither does the reference.
+        if name.endswith(("lora_A.weight", "lora_B.weight")):
+            continue
+        reference.setdefault(name, param.detach().clone())
+
+    # Gradient checkpointing on one projection puts a real
+    # _checkpoint_wrapped_module segment in the module tree — the exact shape
+    # that broke an earlier revision's path stripping.
+    model.k_proj = checkpoint_wrapper(model.k_proj)
+
+    for child in list(model.children()):
+        fully_shard(child)
+    fully_shard(model)
+    return model, reference
+
+
+def _gather(model: nn.Module) -> tuple[dict[str, torch.Tensor], dict[str, str], int]:
+    callback = HFExportCallback(dtype="float32")
+    chunks, manifest, total = callback._gather_weights(SimpleNamespace(model=SimpleNamespace(model=model)))
+    return {k: v for c in chunks for k, v in c.items()}, manifest, total
+
+
+def test_adapters_are_actually_sharded():
+    """Guards the test itself: without DTensors the rest proves nothing."""
+    setup_distributed_environment()
+    model, _ = _build_sharded_block()
+
+    sharded = [n for n, p in model.named_parameters() if isinstance(p, torch.distributed.tensor.DTensor)]
+    assert sharded, "FSDP2 did not shard anything; the remaining assertions would be vacuous"
+    assert len([n for n in sharded if "lora_" in n]) == 8, f"expected 8 sharded adapter tensors, got {sharded}"
+
+
+def test_merged_export_matches_the_unsharded_reference():
+    """The whole contract: sharded export == single-process export, key for key."""
+    rank, _ = setup_distributed_environment()
+    model, reference = _build_sharded_block()
+
+    flat, manifest, total = _gather(model)
+    # Returning on every rank is itself the assertion that the reordered
+    # all-gathers stay in lockstep; a mismatch hangs here instead.
+    dist.barrier()
+
+    if rank != 0:
+        return
+
+    assert not [k for k in flat if "lora_" in k], f"adapter keys leaked into the export: {sorted(flat)}"
+    assert set(flat) == set(reference), (
+        f"extra={sorted(set(flat) - set(reference))} missing={sorted(set(reference) - set(flat))}"
+    )
+    assert set(manifest) == set(flat)
+    assert total == sum(t.element_size() * t.numel() for t in flat.values())
+    for key, tensor in flat.items():
+        torch.testing.assert_close(tensor.cuda(), reference[key], rtol=0, atol=1e-5, msg=f"mismatch at {key}")
+
+
+def test_merge_completeness_guard_fires_under_fsdp():
+    """The guard must abort on every rank, not just where the manifest lives."""
+    setup_distributed_environment()
+    model, _ = _build_sharded_block()
+
+    callback = HFExportCallback(dtype="float32")
+    real_plan = callback._lora_merge_plan
+    # A target that no parameter name can match — i.e. the plan and the loop
+    # disagreeing, which is how a silently unmerged export would arise.
+    callback._lora_merge_plan = lambda root: ({"bogus.path.weight": None}, real_plan(root)[1])
+
+    with pytest.raises(RuntimeError, match="LoRA merge incomplete"):
+        callback._gather_weights(SimpleNamespace(model=SimpleNamespace(model=model)))