From 7a54b498cbff6c8204a89d6c51e0bb4eb152e194 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Fri, 23 Oct 2026 18:54:40 +0000
Subject: [PATCH 1/3] Merge LoRA updates into automatic HF exports without
 changing training weights

Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
---
 .github/workflows/gpu-tests.yml               |   6 +
 cosmos_framework/callbacks/hf_export.py       |  71 ++++-
 cosmos_framework/callbacks/hf_export_test.py  | 247 ++++++++++++++++++
 cosmos_framework/utils/generator/lora.py      |  24 +-
 cosmos_framework/utils/generator/lora_test.py |  30 +++
 5 files changed, 369 insertions(+), 9 deletions(-)
 create mode 100644 cosmos_framework/callbacks/hf_export_test.py

diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml
index ae4e6ce83..426be087e 100644
--- a/.github/workflows/gpu-tests.yml
+++ b/.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
diff --git a/cosmos_framework/callbacks/hf_export.py b/cosmos_framework/callbacks/hf_export.py
index ea6a0e3e5..b4d8af620 100644
--- a/cosmos_framework/callbacks/hf_export.py
+++ b/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,32 @@
 }
 
 
+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]]:
+    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 +238,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 +254,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 +276,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 +308,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
 
+        # Check after all collectives, on every rank. A tied/aliased base weight
+        # omitted by named_parameters() must not silently discard its adapter.
+        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)
diff --git a/cosmos_framework/callbacks/hf_export_test.py b/cosmos_framework/callbacks/hf_export_test.py
new file mode 100644
index 000000000..fb22cb8b8
--- /dev/null
+++ b/cosmos_framework/callbacks/hf_export_test.py
@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Exercise dense HF reloads and the collective path used by checkpoint export."""
+
+import json
+import os
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from torch.distributed.tensor import DTensor
+from transformers import AutoModelForCausalLM, Qwen2Config, Qwen2ForCausalLM
+
+from cosmos_framework.callbacks.hf_export import HFExportCallback
+from cosmos_framework.utils.generator.lora import LoraInjectedLinear, inject_lora_pre_fsdp
+
+pytestmark = pytest.mark.L0
+
+
+def _model(*, lora=True, use_rslora=False, dtype=torch.float32):
+    torch.manual_seed(123)
+    config = Qwen2Config(
+        vocab_size=32,
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        tie_word_embeddings=False,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2ForCausalLM(config).to(dtype=dtype).eval()
+    if lora:
+        inject_lora_pre_fsdp(
+            model,
+            lora_rank=4,
+            lora_alpha=8.0,
+            lora_target_modules="q_proj,v_proj",
+            lora_use_rslora=use_rslora,
+            lora_precision="float32",
+        )
+        with torch.no_grad():
+            for module in model.modules():
+                if isinstance(module, LoraInjectedLinear):
+                    module.lora_A.weight.normal_(std=0.1)
+                    module.lora_B.weight.normal_(std=0.1)
+    return model
+
+
+def _wrapped(model):
+    return SimpleNamespace(model=SimpleNamespace(model=model))
+
+
+def _callback(dtype="float32"):
+    callback = HFExportCallback(dtype=dtype)
+    callback.config = SimpleNamespace(
+        checkpoint=SimpleNamespace(
+            hf_export=SimpleNamespace(upload_to_object_store=SimpleNamespace(enabled=False), hf_repo_id="")
+        )
+    )
+    callback._MAX_SHARD_BYTES = 1024  # Exercise multiple shards and their index.
+    return callback
+
+
+@pytest.mark.parametrize("use_rslora", [False, True])
+@pytest.mark.parametrize(
+    "base_dtype,export_dtype", [(torch.float32, "float32"), (torch.bfloat16, "bfloat16"), (torch.bfloat16, "float32")]
+)
+@torch.no_grad()
+def test_lora_export_reloads_as_dense_hf_model(tmp_path, use_rslora, base_dtype, export_dtype):
+    model = _model(use_rslora=use_rslora, dtype=base_dtype)
+    tokens = torch.tensor([[1, 4, 9, 2]])
+    expected = model(tokens).logits
+    # A zero-initialized adapter would make a broken base-only export pass.
+    base_logits = _model(lora=False, dtype=base_dtype)(tokens).logits
+    assert (expected - base_logits).abs().max() > 0.01
+    before = {name: parameter.clone() for name, parameter in model.named_parameters()}
+    identities = {name: (id(parameter), parameter.requires_grad) for name, parameter in model.named_parameters()}
+    modes = {name: module.training for name, module in model.named_modules()}
+
+    callback = _callback(export_dtype)
+    chunks, manifest, total_size = callback._gather_weights(_wrapped(model))
+    exported = {name: tensor.clone() for chunk in chunks for name, tensor in chunk.items()}
+    assert len(chunks) > 1
+    assert set(exported) == set(_model(lora=False).state_dict())
+    assert total_size == sum(t.numel() * t.element_size() for t in exported.values())
+    callback._do_save_and_upload(chunks, manifest, total_size, model.config, "", str(tmp_path), 1)
+    with (tmp_path / "model.safetensors.index.json").open() as handle:
+        index = json.load(handle)
+    assert index == {"metadata": {"total_size": total_size}, "weight_map": manifest}
+    assert all((tmp_path / shard).is_file() for shard in manifest.values())
+    reloaded, loading = AutoModelForCausalLM.from_pretrained(
+        tmp_path,
+        local_files_only=True,
+        dtype=getattr(torch, export_dtype),
+        attn_implementation="eager",
+        output_loading_info=True,
+    )
+    assert not loading["unexpected_keys"]
+    assert not loading["missing_keys"]
+    assert not loading["mismatched_keys"]
+    tolerance = 2e-3 if base_dtype == torch.bfloat16 else 2e-6
+    torch.testing.assert_close(reloaded(tokens).logits.float(), expected.float(), atol=tolerance, rtol=0.02)
+
+    # Export must not merge the live adapter a second time or change trainability.
+    again, _, _ = callback._gather_weights(_wrapped(model))
+    for chunk in again:
+        for name, tensor in chunk.items():
+            torch.testing.assert_close(tensor, exported[name], rtol=0, atol=0)
+    for name, parameter in model.named_parameters():
+        torch.testing.assert_close(parameter, before[name], rtol=0, atol=0)
+        assert (id(parameter), parameter.requires_grad) == identities[name]
+    assert {name: module.training for name, module in model.named_modules()} == modes
+
+
+@pytest.mark.parametrize("dtype", ["float32", "bfloat16"])
+@torch.no_grad()
+def test_dense_export_preserves_weights_and_reload(tmp_path, dtype):
+    model = _model(lora=False, dtype=getattr(torch, dtype))
+    callback = _callback(dtype)
+    chunks, manifest, total_size = callback._gather_weights(_wrapped(model))
+    for chunk in chunks:
+        for name, tensor in chunk.items():
+            torch.testing.assert_close(tensor, model.state_dict()[name], rtol=0, atol=0)
+    callback._do_save_and_upload(chunks, manifest, total_size, model.config, "", str(tmp_path), 1)
+    reloaded = AutoModelForCausalLM.from_pretrained(
+        tmp_path, local_files_only=True, dtype=getattr(torch, dtype), attn_implementation="eager"
+    )
+    tokens = torch.tensor([[1, 4, 9, 2]])
+    torch.testing.assert_close(reloaded(tokens).logits, model(tokens).logits, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("wrappers", [(), ("_orig_mod",), ("_checkpoint_wrapped_module", "_orig_mod")])
+@torch.no_grad()
+def test_root_lora_and_wrapper_names(wrappers):
+    layer = LoraInjectedLinear(nn.Linear(4, 3), rank=2, alpha=4, use_rslora=True)
+    network = layer
+    for name in wrappers:
+        network = nn.ModuleDict({name: network})
+    chunks, manifest, _ = _callback()._gather_weights(_wrapped(network))
+    exported = {name: tensor for chunk in chunks for name, tensor in chunk.items()}
+    assert set(manifest) == {"weight", "bias"}
+    dense = nn.Linear(4, 3)
+    dense.load_state_dict(exported, strict=True)
+    inputs = torch.randn(2, 4)
+    torch.testing.assert_close(dense(inputs), layer(inputs))
+
+
+@torch.no_grad()
+def test_cpu_export_owns_snapshot_storage():
+    model = nn.Linear(4, 3)
+    chunks, _, _ = _callback()._gather_weights(_wrapped(model))
+    expected = model.weight.clone()
+    model.weight.add_(1)
+    torch.testing.assert_close(chunks[0]["weight"], expected, rtol=0, atol=0)
+
+
+def test_ambiguous_tied_lora_weight_is_rejected():
+    base = nn.Linear(4, 3)
+    network = nn.ModuleDict({"base": base, "adapted": LoraInjectedLinear(base, rank=2, alpha=4)})
+    with pytest.raises(ValueError, match="unmerged LoRA weights=.*adapted.weight"):
+        _callback()._gather_weights(_wrapped(network))
+
+
+def test_wrapper_name_collision_is_rejected():
+    network = nn.ModuleDict({"layer": nn.Linear(4, 3), "_orig_mod": nn.ModuleDict({"layer": nn.Linear(4, 3)})})
+    with pytest.raises(ValueError, match="duplicate keys=.*layer.weight"):
+        _callback()._gather_weights(_wrapped(network))
+
+
+def test_similar_names_are_not_removed():
+    network = nn.ModuleDict({"my_orig_mod": nn.Linear(4, 3), "lora_A": nn.Linear(4, 3)})
+    _, manifest, _ = _callback()._gather_weights(_wrapped(network))
+    assert set(manifest) == set(network.state_dict())
+
+
+@pytest.fixture(scope="module")
+def fsdp_group():
+    """Keep the process group alive across cases, including DTensor's cached ops."""
+    if os.environ.get("WORLD_SIZE") != "2":
+        pytest.skip("requires torchrun --nproc_per_node=2 for actual FSDP/NCCL collectives")
+    rank = int(os.environ["RANK"])
+    device = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(device)
+    torch.distributed.init_process_group("nccl", timeout=timedelta(seconds=60))
+    try:
+        yield rank, device
+    finally:
+        torch.distributed.destroy_process_group()
+
+
+@pytest.mark.GPU
+@pytest.mark.parametrize("lora", [False, True])
+def test_fsdp_export_gathers_once_in_rank_order_and_reloads(tmp_path, lora, fsdp_group):
+    """Run with torchrun --nproc_per_node=2, like the other distributed CI cases."""
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import fully_shard
+
+    rank, device = fsdp_group
+    output = str(tmp_path / "export")
+    model = _model(lora=lora, use_rslora=True).cuda()
+    tokens = torch.tensor([[1, 4, 9, 2]], device=device)
+    with torch.no_grad():
+        expected = model(tokens).logits.cpu()
+    fully_shard(model, mesh=init_device_mesh("cuda", (2,)))
+    parameters = dict(model.named_parameters())
+    assert all(isinstance(parameter, DTensor) for parameter in parameters.values())
+    before = {name: parameter.full_tensor().detach().clone() for name, parameter in parameters.items()}
+    names = {id(parameter): name for name, parameter in parameters.items()}
+    order = []
+    full_tensor = DTensor.full_tensor
+
+    def recording_full_tensor(tensor, *args, **kwargs):
+        order.append(names[id(tensor)])
+        return full_tensor(tensor, *args, **kwargs)
+
+    callback = _callback()
+    DTensor.full_tensor = recording_full_tensor
+    try:
+        chunks, manifest, total_size = callback._gather_weights(_wrapped(model))
+    finally:
+        DTensor.full_tensor = full_tensor
+    orders = [None, None]
+    torch.distributed.all_gather_object(orders, order)
+    assert orders[0] == orders[1]
+    assert len(order) == len(parameters) and set(order) == set(parameters)
+    for name, parameter in model.named_parameters():
+        assert parameter is parameters[name]
+        torch.testing.assert_close(parameter.full_tensor(), before[name], rtol=0, atol=0)
+    if rank == 0:
+        callback._do_save_and_upload(chunks, manifest, total_size, model.config, "", output, 1)
+        dense, info = AutoModelForCausalLM.from_pretrained(
+            output,
+            local_files_only=True,
+            dtype=torch.float32,
+            attn_implementation="eager",
+            output_loading_info=True,
+        )
+        assert not info["unexpected_keys"] and not info["missing_keys"]
+        with torch.no_grad():
+            torch.testing.assert_close(dense(tokens.cpu()).logits, expected, atol=2e-6, rtol=2e-5)
+    else:
+        assert (chunks, manifest, total_size) == ([], {}, 0)
+    torch.distributed.barrier()
diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py
index 2cdca0569..e588afb3c 100644
--- a/cosmos_framework/utils/generator/lora.py
+++ b/cosmos_framework/utils/generator/lora.py
@@ -92,11 +92,31 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         lora_out = self.lora_B(self.lora_A(adapter_input)).to(dtype=base_out.dtype)
         return base_out + self._lora_scale * lora_out
 
+    @torch.no_grad()
+    def merged_weight(
+        self,
+        base: torch.Tensor,
+        lora_a: torch.Tensor,
+        lora_b: torch.Tensor,
+        *,
+        dtype: torch.dtype | None = None,
+    ) -> torch.Tensor:
+        """Return a dense weight without changing the base or adapter parameters.
+
+        Callers exporting sharded models must gather all three tensors first.
+        Accumulate in at least FP32 and cast only the completed merge, including
+        when FP32 adapters are trained alongside a BF16 base. Both offline and
+        checkpoint-time exports use the module's LoRA/RSLoRA scale here.
+        """
+        accumulation_dtype = torch.float64 if base.dtype == torch.float64 else torch.float32
+        with torch.autocast(device_type=base.device.type, enabled=False):
+            delta = lora_b.to(accumulation_dtype) @ lora_a.to(accumulation_dtype)
+            return (base.to(accumulation_dtype) + self._lora_scale * delta).to(dtype=dtype or base.dtype)
+
     @torch.no_grad()
     def merge_adapter_(self) -> None:
         """Merge the adapter update into the base weight in place."""
-        delta = self.lora_B.weight.float() @ self.lora_A.weight.float()
-        self.weight.add_(delta.to(dtype=self.weight.dtype), alpha=self._lora_scale)
+        self.weight.copy_(self.merged_weight(self.weight, self.lora_A.weight, self.lora_B.weight))
 
 
 def _target_matches(full_child_path: str, child_name: str, target: str) -> bool:
diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py
index 5af17aba8..0baa1e4d7 100644
--- a/cosmos_framework/utils/generator/lora_test.py
+++ b/cosmos_framework/utils/generator/lora_test.py
@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import pytest
 import torch
 
 from cosmos_framework.utils.generator.lora import (
@@ -12,6 +13,8 @@
     merge_lora_adapters_,
 )
 
+pytestmark = pytest.mark.L0
+
 
 class _Tiny(torch.nn.Module):
     def __init__(self) -> None:
@@ -81,3 +84,30 @@ def test_rslora_merge_matches_unmerged_forward() -> None:
         model.q_proj.lora_A.weight.zero_()
         model.q_proj.lora_B.weight.zero_()
     torch.testing.assert_close(model.q_proj(values), expected)
+
+
+@pytest.mark.parametrize("use_rslora", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float64])
+def test_export_merge_and_inplace_merge_share_scaling_and_precision(dtype, use_rslora):
+    torch.manual_seed(13)
+    layer = LoraInjectedLinear(
+        torch.nn.Linear(4, 3, dtype=dtype), rank=2, alpha=3, use_rslora=use_rslora, adapter_dtype=torch.float32
+    )
+    base = layer.weight.detach().clone()
+    adapter_a = layer.lora_A.weight.detach().clone()
+    adapter_b = layer.lora_B.weight.detach().clone()
+    scale = 3 / (2**0.5 if use_rslora else 2)
+    accumulation_dtype = torch.float64 if dtype == torch.float64 else torch.float32
+    expected = base.to(accumulation_dtype) + scale * (
+        adapter_b.to(accumulation_dtype) @ adapter_a.to(accumulation_dtype)
+    )
+    # A higher-precision export must not first round the merged weight to BF16.
+    with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
+        exported = layer.merged_weight(base, adapter_a, adapter_b, dtype=accumulation_dtype)
+    torch.testing.assert_close(exported, expected, rtol=0, atol=0)
+    for actual, original in ((layer.weight, base), (layer.lora_A.weight, adapter_a), (layer.lora_B.weight, adapter_b)):
+        torch.testing.assert_close(actual, original, rtol=0, atol=0)
+    weight_identity = id(layer.weight)
+    layer.merge_adapter_()
+    assert id(layer.weight) == weight_identity
+    torch.testing.assert_close(layer.weight, expected.to(dtype), rtol=0, atol=0)

From 28e15bea41653cb54cc67e631b2e073ff03b8272 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Sat, 24 Oct 2026 15:53:17 +0000
Subject: [PATCH 2/3] Validate shared LoRA export aliases and preserve legacy
 merge behavior

Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
---
 cosmos_framework/callbacks/hf_export.py       |  28 ++++-
 cosmos_framework/callbacks/hf_export_test.py  | 102 +++++++++++++++---
 cosmos_framework/utils/generator/lora.py      |   8 +-
 cosmos_framework/utils/generator/lora_test.py |  45 +++++++-
 4 files changed, 161 insertions(+), 22 deletions(-)

diff --git a/cosmos_framework/callbacks/hf_export.py b/cosmos_framework/callbacks/hf_export.py
index b4d8af620..e754f42c8 100644
--- a/cosmos_framework/callbacks/hf_export.py
+++ b/cosmos_framework/callbacks/hf_export.py
@@ -60,6 +60,30 @@ def _hf_parameter_name(name: str) -> str:
 
 
 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():
@@ -315,8 +339,8 @@ def _gather_weights(self, model: Any) -> tuple[list[dict[str, torch.Tensor]], di
                 total_size += tensor_bytes
             current_chunk_bytes += tensor_bytes
 
-        # Check after all collectives, on every rank. A tied/aliased base weight
-        # omitted by named_parameters() must not silently discard its adapter.
+        # 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(
diff --git a/cosmos_framework/callbacks/hf_export_test.py b/cosmos_framework/callbacks/hf_export_test.py
index fb22cb8b8..780cbac04 100644
--- a/cosmos_framework/callbacks/hf_export_test.py
+++ b/cosmos_framework/callbacks/hf_export_test.py
@@ -20,7 +20,17 @@
 pytestmark = pytest.mark.L0
 
 
-def _model(*, lora=True, use_rslora=False, dtype=torch.float32):
+class _SharedBaseNetwork(nn.Module):
+    def __init__(self, layers):
+        super().__init__()
+        for name, layer in layers:
+            self.add_module(name, layer)
+
+    def forward(self, inputs):
+        return self.base(inputs) + self.adapted(inputs)
+
+
+def _model(*, lora=True, use_rslora=False, dtype=torch.float32, tied_embeddings=False):
     torch.manual_seed(123)
     config = Qwen2Config(
         vocab_size=32,
@@ -29,7 +39,7 @@ def _model(*, lora=True, use_rslora=False, dtype=torch.float32):
         num_hidden_layers=1,
         num_attention_heads=2,
         num_key_value_heads=2,
-        tie_word_embeddings=False,
+        tie_word_embeddings=tied_embeddings,
     )
     config._attn_implementation = "eager"
     model = Qwen2ForCausalLM(config).to(dtype=dtype).eval()
@@ -117,18 +127,26 @@ def test_lora_export_reloads_as_dense_hf_model(tmp_path, use_rslora, base_dtype,
 
 
 @pytest.mark.parametrize("dtype", ["float32", "bfloat16"])
+@pytest.mark.parametrize("tied_embeddings", [False, True])
 @torch.no_grad()
-def test_dense_export_preserves_weights_and_reload(tmp_path, dtype):
-    model = _model(lora=False, dtype=getattr(torch, dtype))
+def test_dense_export_preserves_weights_and_reload(tmp_path, dtype, tied_embeddings):
+    model = _model(lora=False, dtype=getattr(torch, dtype), tied_embeddings=tied_embeddings)
     callback = _callback(dtype)
     chunks, manifest, total_size = callback._gather_weights(_wrapped(model))
     for chunk in chunks:
         for name, tensor in chunk.items():
             torch.testing.assert_close(tensor, model.state_dict()[name], rtol=0, atol=0)
     callback._do_save_and_upload(chunks, manifest, total_size, model.config, "", str(tmp_path), 1)
-    reloaded = AutoModelForCausalLM.from_pretrained(
-        tmp_path, local_files_only=True, dtype=getattr(torch, dtype), attn_implementation="eager"
+    reloaded, loading = AutoModelForCausalLM.from_pretrained(
+        tmp_path,
+        local_files_only=True,
+        dtype=getattr(torch, dtype),
+        attn_implementation="eager",
+        output_loading_info=True,
     )
+    assert not loading["unexpected_keys"] and not loading["missing_keys"] and not loading["mismatched_keys"]
+    if tied_embeddings:
+        assert reloaded.get_input_embeddings().weight is reloaded.get_output_embeddings().weight
     tokens = torch.tensor([[1, 4, 9, 2]])
     torch.testing.assert_close(reloaded(tokens).logits, model(tokens).logits, rtol=0, atol=0)
 
@@ -158,13 +176,42 @@ def test_cpu_export_owns_snapshot_storage():
     torch.testing.assert_close(chunks[0]["weight"], expected, rtol=0, atol=0)
 
 
-def test_ambiguous_tied_lora_weight_is_rejected():
+@pytest.mark.parametrize("adapted_first", [False, True])
+def test_ambiguous_tied_lora_weight_is_rejected(adapted_first):
+    base = nn.Linear(4, 3)
+    adapted = LoraInjectedLinear(base, rank=2, alpha=4)
+    nn.init.constant_(adapted.lora_A.weight, 0.5)
+    nn.init.constant_(adapted.lora_B.weight, 0.25)
+    layers = [("base", base), ("adapted", adapted)]
+    network = nn.ModuleDict(reversed(layers) if adapted_first else layers)
+    with pytest.raises(ValueError, match="LoRA"):
+        _callback()._gather_weights(_wrapped(network))
+
+
+@pytest.mark.parametrize("reverse", [False, True])
+def test_distinct_adapters_sharing_a_base_are_rejected(reverse):
     base = nn.Linear(4, 3)
-    network = nn.ModuleDict({"base": base, "adapted": LoraInjectedLinear(base, rank=2, alpha=4)})
-    with pytest.raises(ValueError, match="unmerged LoRA weights=.*adapted.weight"):
+    layers = [(name, LoraInjectedLinear(base, rank=2, alpha=alpha)) for name, alpha in (("first", 2), ("second", 4))]
+    network = nn.ModuleDict(reversed(layers) if reverse else layers)
+    with pytest.raises(ValueError, match="LoRA"):
         _callback()._gather_weights(_wrapped(network))
 
 
+@torch.no_grad()
+def test_aliases_of_the_same_lora_module_export_once():
+    layer = LoraInjectedLinear(nn.Linear(4, 3), rank=2, alpha=4)
+    nn.init.constant_(layer.lora_A.weight, 0.5)
+    nn.init.constant_(layer.lora_B.weight, 0.25)
+    network = nn.ModuleDict({"first": layer, "second": layer})
+    chunks, manifest, _ = _callback()._gather_weights(_wrapped(network))
+    assert set(manifest) == {"first.weight", "first.bias"}
+    dense = nn.Linear(4, 3)
+    dense.load_state_dict({name.removeprefix("first."): tensor for chunk in chunks for name, tensor in chunk.items()})
+    inputs = torch.randn(2, 4)
+    torch.testing.assert_close(dense(inputs), network["first"](inputs))
+    torch.testing.assert_close(dense(inputs), network["second"](inputs))
+
+
 def test_wrapper_name_collision_is_rejected():
     network = nn.ModuleDict({"layer": nn.Linear(4, 3), "_orig_mod": nn.ModuleDict({"layer": nn.Linear(4, 3)})})
     with pytest.raises(ValueError, match="duplicate keys=.*layer.weight"):
@@ -193,15 +240,15 @@ def fsdp_group():
 
 
 @pytest.mark.GPU
-@pytest.mark.parametrize("lora", [False, True])
-def test_fsdp_export_gathers_once_in_rank_order_and_reloads(tmp_path, lora, fsdp_group):
+@pytest.mark.parametrize("lora,tied_embeddings", [(False, False), (True, False), (False, True)])
+def test_fsdp_export_gathers_once_in_rank_order_and_reloads(tmp_path, lora, tied_embeddings, fsdp_group):
     """Run with torchrun --nproc_per_node=2, like the other distributed CI cases."""
     from torch.distributed.device_mesh import init_device_mesh
     from torch.distributed.fsdp import fully_shard
 
     rank, device = fsdp_group
     output = str(tmp_path / "export")
-    model = _model(lora=lora, use_rslora=True).cuda()
+    model = _model(lora=lora, use_rslora=True, tied_embeddings=tied_embeddings).cuda()
     tokens = torch.tensor([[1, 4, 9, 2]], device=device)
     with torch.no_grad():
         expected = model(tokens).logits.cpu()
@@ -240,8 +287,39 @@ def recording_full_tensor(tensor, *args, **kwargs):
             output_loading_info=True,
         )
         assert not info["unexpected_keys"] and not info["missing_keys"]
+        if tied_embeddings:
+            assert dense.get_input_embeddings().weight is dense.get_output_embeddings().weight
         with torch.no_grad():
             torch.testing.assert_close(dense(tokens.cpu()).logits, expected, atol=2e-6, rtol=2e-5)
     else:
         assert (chunks, manifest, total_size) == ([], {}, 0)
     torch.distributed.barrier()
+
+
+@pytest.mark.GPU
+@pytest.mark.parametrize("adapted_first", [False, True])
+def test_fsdp_rejects_shared_lora_base_before_collectives(adapted_first, fsdp_group, monkeypatch):
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import fully_shard
+
+    _, device = fsdp_group
+    base = nn.Linear(4, 3, bias=False, device=device)
+    adapted = LoraInjectedLinear(base, rank=2, alpha=4)
+    layers = [("base", base), ("adapted", adapted)]
+    network = _SharedBaseNetwork(reversed(layers) if adapted_first else layers)
+    fully_shard(network, mesh=init_device_mesh("cuda", (2,)))
+    assert network.base.weight is network.adapted.weight
+    assert isinstance(base.weight, DTensor)
+    before = base.weight.to_local().detach().clone()
+
+    def unexpected_gather(*args, **kwargs):
+        pytest.fail("Ambiguous aliases must be rejected before any tensor collective")
+
+    monkeypatch.setattr(DTensor, "full_tensor", unexpected_gather)
+    with pytest.raises(ValueError, match="LoRA") as error:
+        _callback()._gather_weights(_wrapped(network))
+    messages = [None, None]
+    torch.distributed.all_gather_object(messages, str(error.value))
+    assert messages[0] == messages[1]
+    assert "adapted.weight" in messages[0] and "base.weight" in messages[0]
+    torch.testing.assert_close(base.weight.to_local(), before, rtol=0, atol=0)
diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py
index e588afb3c..c0206d3cb 100644
--- a/cosmos_framework/utils/generator/lora.py
+++ b/cosmos_framework/utils/generator/lora.py
@@ -105,8 +105,9 @@ def merged_weight(
 
         Callers exporting sharded models must gather all three tensors first.
         Accumulate in at least FP32 and cast only the completed merge, including
-        when FP32 adapters are trained alongside a BF16 base. Both offline and
-        checkpoint-time exports use the module's LoRA/RSLoRA scale here.
+        when FP32 adapters are trained alongside a BF16 base. This is the
+        automatic-export policy; the existing in-place/offline merge retains
+        its delta rounding and ambient autocast behavior in ``merge_adapter_``.
         """
         accumulation_dtype = torch.float64 if base.dtype == torch.float64 else torch.float32
         with torch.autocast(device_type=base.device.type, enabled=False):
@@ -116,7 +117,8 @@ def merged_weight(
     @torch.no_grad()
     def merge_adapter_(self) -> None:
         """Merge the adapter update into the base weight in place."""
-        self.weight.copy_(self.merged_weight(self.weight, self.lora_A.weight, self.lora_B.weight))
+        delta = self.lora_B.weight.float() @ self.lora_A.weight.float()
+        self.weight.add_(delta.to(dtype=self.weight.dtype), alpha=self._lora_scale)
 
 
 def _target_matches(full_child_path: str, child_name: str, target: str) -> bool:
diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py
index 0baa1e4d7..7242c57d9 100644
--- a/cosmos_framework/utils/generator/lora_test.py
+++ b/cosmos_framework/utils/generator/lora_test.py
@@ -10,6 +10,7 @@
     LoraInjectedLinear,
     apply_lora_trainable_scope,
     inject_lora_pre_fsdp,
+    merge_and_strip_lora_adapters_,
     merge_lora_adapters_,
 )
 
@@ -88,7 +89,7 @@ def test_rslora_merge_matches_unmerged_forward() -> None:
 
 @pytest.mark.parametrize("use_rslora", [False, True])
 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.float64])
-def test_export_merge_and_inplace_merge_share_scaling_and_precision(dtype, use_rslora):
+def test_export_merge_preserves_precision_and_training_weights(dtype, use_rslora):
     torch.manual_seed(13)
     layer = LoraInjectedLinear(
         torch.nn.Linear(4, 3, dtype=dtype), rank=2, alpha=3, use_rslora=use_rslora, adapter_dtype=torch.float32
@@ -107,7 +108,41 @@ def test_export_merge_and_inplace_merge_share_scaling_and_precision(dtype, use_r
     torch.testing.assert_close(exported, expected, rtol=0, atol=0)
     for actual, original in ((layer.weight, base), (layer.lora_A.weight, adapter_a), (layer.lora_B.weight, adapter_b)):
         torch.testing.assert_close(actual, original, rtol=0, atol=0)
-    weight_identity = id(layer.weight)
-    layer.merge_adapter_()
-    assert id(layer.weight) == weight_identity
-    torch.testing.assert_close(layer.weight, expected.to(dtype), rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("merge_api", ["layer", "network", "offline"])
+@pytest.mark.parametrize(
+    "dtype,autocast,expected_inplace,expected_export",
+    [
+        (torch.bfloat16, False, 1.0, 1.0078125),
+        (torch.bfloat16, True, 1.0, 1.0078125),
+        (torch.float32, False, 1.00392, 1.00392),
+        (torch.float32, True, 1.00390625, 1.00392),
+    ],
+)
+@torch.no_grad()
+def test_existing_merge_apis_retain_rounding_and_autocast(
+    merge_api, dtype, autocast, expected_inplace, expected_export
+):
+    layer = LoraInjectedLinear(
+        torch.nn.Linear(1, 1, bias=False, dtype=dtype), rank=1, alpha=1, adapter_dtype=torch.float32
+    )
+    layer.weight.fill_(1.0)
+    layer.lora_A.weight.fill_(1.0)
+    layer.lora_B.weight.fill_(0.00392)
+    weight = layer.weight
+    network = torch.nn.ModuleDict({"layer": layer})
+    with torch.autocast(device_type="cpu", dtype=torch.bfloat16, enabled=autocast):
+        exported = layer.merged_weight(layer.weight, layer.lora_A.weight, layer.lora_B.weight)
+        if merge_api == "layer":
+            layer.merge_adapter_()
+        elif merge_api == "network":
+            assert merge_lora_adapters_(network) == 1
+        else:
+            assert merge_and_strip_lora_adapters_(network) == 1
+            assert type(network["layer"]) is torch.nn.Linear
+    # Deliberate boundary: legacy APIs round the delta before adding it, and
+    # honor ambient autocast. Automatic export rounds only the completed sum.
+    assert network["layer"].weight is weight
+    torch.testing.assert_close(weight, torch.full_like(weight, expected_inplace), rtol=0, atol=0)
+    torch.testing.assert_close(exported, torch.full_like(weight, expected_export), rtol=0, atol=0)

From 0b43ed1d673cb012c75d12066decc85613787f65 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Mon, 26 Oct 2026 20:22:21 +0000
Subject: [PATCH 3/3] Preserve FP32 precision in automatic LoRA exports

Disable TF32 during the non-mutating CUDA merge and restore the caller's
matmul precision on success and failure. Keep the existing in-place and
offline merge APIs unchanged.

Compare CPU HF reloads against CPU references, exercise export with TF32
on and off, and propagate rank-local validation errors before peers move
to another distributed test. Cover LoRA/RSLoRA precision and setting
restoration, including merge exceptions.

Signed-off-by: Ramanathan Arunachalam <rarunachalam@nvidia.com>
---
 cosmos_framework/callbacks/hf_export_test.py  | 86 ++++++++++++++-----
 cosmos_framework/utils/generator/lora.py      | 16 +++-
 cosmos_framework/utils/generator/lora_test.py | 26 ++++++
 3 files changed, 103 insertions(+), 25 deletions(-)

diff --git a/cosmos_framework/callbacks/hf_export_test.py b/cosmos_framework/callbacks/hf_export_test.py
index 780cbac04..3980e1714 100644
--- a/cosmos_framework/callbacks/hf_export_test.py
+++ b/cosmos_framework/callbacks/hf_export_test.py
@@ -5,6 +5,8 @@
 
 import json
 import os
+import traceback
+from contextlib import contextmanager
 from datetime import timedelta
 from types import SimpleNamespace
 
@@ -241,17 +243,22 @@ def fsdp_group():
 
 @pytest.mark.GPU
 @pytest.mark.parametrize("lora,tied_embeddings", [(False, False), (True, False), (False, True)])
-def test_fsdp_export_gathers_once_in_rank_order_and_reloads(tmp_path, lora, tied_embeddings, fsdp_group):
+def test_fsdp_export_gathers_once_in_rank_order_and_reloads(
+    tmp_path, lora, tied_embeddings, fsdp_group, export_matmul_precision
+):
     """Run with torchrun --nproc_per_node=2, like the other distributed CI cases."""
     from torch.distributed.device_mesh import init_device_mesh
     from torch.distributed.fsdp import fully_shard
 
-    rank, device = fsdp_group
+    rank, _ = fsdp_group
     output = str(tmp_path / "export")
-    model = _model(lora=lora, use_rslora=True, tied_embeddings=tied_embeddings).cuda()
-    tokens = torch.tensor([[1, 4, 9, 2]], device=device)
+    model = _model(lora=lora, use_rslora=True, tied_embeddings=tied_embeddings)
+    tokens = torch.tensor([[1, 4, 9, 2]])
+    # The reload runs on CPU, so compute its reference on CPU too. CUDA may use
+    # TF32 kernels, whose rounding is unrelated to checkpoint correctness.
     with torch.no_grad():
-        expected = model(tokens).logits.cpu()
+        expected = model(tokens).logits
+    model = model.cuda()
     fully_shard(model, mesh=init_device_mesh("cuda", (2,)))
     parameters = dict(model.named_parameters())
     assert all(isinstance(parameter, DTensor) for parameter in parameters.values())
@@ -277,22 +284,59 @@ def recording_full_tensor(tensor, *args, **kwargs):
     for name, parameter in model.named_parameters():
         assert parameter is parameters[name]
         torch.testing.assert_close(parameter.full_tensor(), before[name], rtol=0, atol=0)
-    if rank == 0:
-        callback._do_save_and_upload(chunks, manifest, total_size, model.config, "", output, 1)
-        dense, info = AutoModelForCausalLM.from_pretrained(
-            output,
-            local_files_only=True,
-            dtype=torch.float32,
-            attn_implementation="eager",
-            output_loading_info=True,
-        )
-        assert not info["unexpected_keys"] and not info["missing_keys"]
-        if tied_embeddings:
-            assert dense.get_input_embeddings().weight is dense.get_output_embeddings().weight
-        with torch.no_grad():
-            torch.testing.assert_close(dense(tokens.cpu()).logits, expected, atol=2e-6, rtol=2e-5)
-    else:
-        assert (chunks, manifest, total_size) == ([], {}, 0)
+    with _collective_export_check():
+        if rank == 0:
+            callback._do_save_and_upload(chunks, manifest, total_size, model.config, "", output, 1)
+            dense, info = AutoModelForCausalLM.from_pretrained(
+                output,
+                local_files_only=True,
+                dtype=torch.float32,
+                attn_implementation="eager",
+                output_loading_info=True,
+            )
+            assert not info["unexpected_keys"] and not info["missing_keys"]
+            if tied_embeddings:
+                assert dense.get_input_embeddings().weight is dense.get_output_embeddings().weight
+            with torch.no_grad():
+                torch.testing.assert_close(dense(tokens).logits, expected, atol=2e-6, rtol=2e-5)
+        else:
+            assert (chunks, manifest, total_size) == ([], {}, 0)
+
+
+@pytest.fixture(params=["highest", "high"], ids=["fp32", "tf32"])
+def export_matmul_precision(request):
+    previous = torch.get_float32_matmul_precision()
+    torch.set_float32_matmul_precision(request.param)
+    try:
+        yield
+    finally:
+        torch.set_float32_matmul_precision(previous)
+
+
+@contextmanager
+def _collective_export_check():
+    """Report local validation errors on every rank before starting another case."""
+    error = None
+    try:
+        yield
+    except Exception:
+        error = traceback.format_exc()
+    errors = [None] * torch.distributed.get_world_size()
+    torch.distributed.all_gather_object(errors, error)
+    failures = "\n".join(f"Rank {rank}:\n{error}" for rank, error in enumerate(errors) if error is not None)
+    assert not failures, failures
+
+
+@pytest.mark.GPU
+@pytest.mark.parametrize("failing_rank,error_type", [(0, AssertionError), (1, OSError)])
+def test_fsdp_export_validation_failure_reaches_all_ranks(fsdp_group, failing_rank, error_type):
+    rank, _ = fsdp_group
+    with pytest.raises(AssertionError, match=rf"(?s)Rank {failing_rank}:.*injected export validation failure"):
+        with _collective_export_check():
+            if rank == failing_rank:
+                raise error_type("injected export validation failure")
+    # The failure must leave the group usable for the next case, not strand a
+    # peer in the old barrier until the NCCL watchdog kills the test process.
     torch.distributed.barrier()
 
 
diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py
index c0206d3cb..001060ad6 100644
--- a/cosmos_framework/utils/generator/lora.py
+++ b/cosmos_framework/utils/generator/lora.py
@@ -105,14 +105,22 @@ def merged_weight(
 
         Callers exporting sharded models must gather all three tensors first.
         Accumulate in at least FP32 and cast only the completed merge, including
-        when FP32 adapters are trained alongside a BF16 base. This is the
+        when FP32 adapters are trained alongside a BF16 base. Disable TF32 as
+        well as autocast so CUDA matmuls preserve FP32 input precision. This is the
         automatic-export policy; the existing in-place/offline merge retains
         its delta rounding and ambient autocast behavior in ``merge_adapter_``.
         """
         accumulation_dtype = torch.float64 if base.dtype == torch.float64 else torch.float32
-        with torch.autocast(device_type=base.device.type, enabled=False):
-            delta = lora_b.to(accumulation_dtype) @ lora_a.to(accumulation_dtype)
-            return (base.to(accumulation_dtype) + self._lora_scale * delta).to(dtype=dtype or base.dtype)
+        previous_precision = torch.get_float32_matmul_precision() if base.is_cuda else None
+        try:
+            if previous_precision is not None:
+                torch.set_float32_matmul_precision("highest")
+            with torch.autocast(device_type=base.device.type, enabled=False):
+                delta = lora_b.to(accumulation_dtype) @ lora_a.to(accumulation_dtype)
+                return (base.to(accumulation_dtype) + self._lora_scale * delta).to(dtype=dtype or base.dtype)
+        finally:
+            if previous_precision is not None:
+                torch.set_float32_matmul_precision(previous_precision)
 
     @torch.no_grad()
     def merge_adapter_(self) -> None:
diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py
index 7242c57d9..983c3bb89 100644
--- a/cosmos_framework/utils/generator/lora_test.py
+++ b/cosmos_framework/utils/generator/lora_test.py
@@ -146,3 +146,29 @@ def test_existing_merge_apis_retain_rounding_and_autocast(
     assert network["layer"].weight is weight
     torch.testing.assert_close(weight, torch.full_like(weight, expected_inplace), rtol=0, atol=0)
     torch.testing.assert_close(exported, torch.full_like(weight, expected_export), rtol=0, atol=0)
+
+
+@pytest.mark.GPU
+@pytest.mark.parametrize("precision", ["highest", "high", "medium"])
+@pytest.mark.parametrize("use_rslora", [False, True])
+@torch.no_grad()
+def test_cuda_export_merge_preserves_precision_and_restores_matmul_settings(precision, use_rslora):
+    torch.manual_seed(13)
+    layer = LoraInjectedLinear(torch.nn.Linear(16, 16), rank=4, alpha=8, use_rslora=use_rslora)
+    layer.lora_A.weight.normal_(std=0.1)
+    layer.lora_B.weight.normal_(std=0.1)
+    expected = (
+        layer.weight.double() + layer._lora_scale * (layer.lora_B.weight.double() @ layer.lora_A.weight.double())
+    ).float()
+    layer.cuda()
+    previous = torch.get_float32_matmul_precision()
+    torch.set_float32_matmul_precision(precision)
+    try:
+        actual = layer.merged_weight(layer.weight, layer.lora_A.weight, layer.lora_B.weight)
+        assert torch.get_float32_matmul_precision() == precision
+        torch.testing.assert_close(actual.cpu(), expected, atol=1e-7, rtol=1e-6)
+        with pytest.raises(RuntimeError, match="shapes cannot be multiplied"):
+            layer.merged_weight(layer.weight, layer.lora_A.weight, layer.lora_B.weight[:, :-1])
+        assert torch.get_float32_matmul_precision() == precision
+    finally:
+        torch.set_float32_matmul_precision(previous)