diff --git a/.github/workflows/gpu-tests.yml b/.github/workflows/gpu-tests.yml index ae4e6ce8..426be087 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 ea6a0e3e..e754f42c 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,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) diff --git a/cosmos_framework/callbacks/hf_export_test.py b/cosmos_framework/callbacks/hf_export_test.py new file mode 100644 index 00000000..3980e171 --- /dev/null +++ b/cosmos_framework/callbacks/hf_export_test.py @@ -0,0 +1,369 @@ +# 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 +import traceback +from contextlib import contextmanager +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 + + +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, + hidden_size=16, + intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + tie_word_embeddings=tied_embeddings, + ) + 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"]) +@pytest.mark.parametrize("tied_embeddings", [False, True]) +@torch.no_grad() +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, 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) + + +@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) + + +@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) + 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"): + _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,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, 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, _ = fsdp_group + output = str(tmp_path / "export") + 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 + 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()) + 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) + 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() + + +@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 2cdca056..001060ad 100644 --- a/cosmos_framework/utils/generator/lora.py +++ b/cosmos_framework/utils/generator/lora.py @@ -92,6 +92,36 @@ 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. 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 + 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: """Merge the adapter update into the base weight in place.""" diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py index 5af17aba..983c3bb8 100644 --- a/cosmos_framework/utils/generator/lora_test.py +++ b/cosmos_framework/utils/generator/lora_test.py @@ -3,15 +3,19 @@ from __future__ import annotations +import pytest import torch from cosmos_framework.utils.generator.lora import ( LoraInjectedLinear, apply_lora_trainable_scope, inject_lora_pre_fsdp, + merge_and_strip_lora_adapters_, merge_lora_adapters_, ) +pytestmark = pytest.mark.L0 + class _Tiny(torch.nn.Module): def __init__(self) -> None: @@ -81,3 +85,90 @@ 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_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 + ) + 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) + + +@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) + + +@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)