[bugfix] Restore CPU-backed NVFP4 weights for LoRA lifecycle updates (#2076)

Co-authored-by: Yixing Wang <269992774+leo2ultimate@users.noreply.github.com>
Co-authored-by: leo <yixingwang@YIXINGs-MacBook-Pro.local>

docs/inference/h3_community_lora.md

@@ -67,9 +67,9 @@ The FastH3 example scripts take the same options. `basic_fasth3_lora_preview.py`
 `--lora-strength`, and `examples/serving/openai_fasth3_lora.yaml` shows the YAML form.
 
 `lora_strength` multiplies the adapter delta. `1.0` is the published scale and `0` is exactly the model without the
-adapter. Values above `1.0` exaggerate the adapter and usually degrade quality. After construction, you can change
-the strength with `generator.set_lora_adapter(nickname, strength=...)` only when the adapter runs as a side branch, described
-next. Otherwise, create a new generator.
+adapter. Values above `1.0` exaggerate the adapter and usually degrade quality. For a pure low-rank adapter,
+`generator.set_lora_adapter(nickname, strength=...)` changes the strength at runtime. Adapters with dense
+replacement or additive weights require a new generator to change their strength.
 
 ## LoRAs trained on base H3
 
@@ -94,11 +94,13 @@ The extra cost is two small matmuls per adapted linear.
 
 A dense bf16 checkpoint that FastVideo quantizes at load time (for example `transformer_quant: NVFP4` on a bf16
 export) still merges the adapter first and quantizes afterwards, so the quantized weight includes the adapter.
-If quantization retains the original dense weight, adapters continue to merge into that weight and FastVideo
-rebuilds the packed buffers after each adapter change.
+If the NVFP4 layer releases its working weight after packing, its merge-mode LoRA wrapper can restore it from
+the original unsharded CPU snapshot. Changing a pure low-rank adapter rebuilds the packed weights and scales.
+Accumulating another adapter first reapplies the previous contributions in order. Unmerging multiple accumulated
+adapters currently drops earlier contributions from inference ([#2096](https://github.com/hao-ai-lab/FastVideo/issues/2096)).
 
-To make the adapter strength changeable at runtime on a dense checkpoint, set `FASTVIDEO_LORA_SIDE_BRANCH=1`. The
-adapter then runs as a side branch for every layer, at the cost of the extra matmuls. See
+To change strength on a dense checkpoint without remerging, set `FASTVIDEO_LORA_SIDE_BRANCH=1`. The
+adapter then runs as a side branch on supported layers, at the cost of the extra matmuls. See
 [Environment variables](../contributing/env_vars.md).
 
 ## When the whole model does not fit

fastvideo/hooks/layerwise_offload.py

@@ -23,7 +23,7 @@ _BUFFER_OFFLOAD_MIN_BYTES = 1 << 20
 
 
 def _offload_tensors(module: nn.Module, names: dict[str, torch.Tensor] | None = None):
-    """``(name, tensor)`` for every parameter and, when enabled, every large buffer.
+    """``(name, tensor)`` for parameters, placement markers and optionally large buffers.
 
     ``names`` restricts the walk to the tensors chosen at init: an offloaded buffer is a
     zero-element placeholder afterwards and would fail the size test.
@@ -34,10 +34,12 @@ def _offload_tensors(module: nn.Module, names: dict[str, torch.Tensor] | None =
                 yield name, tensor
         return
     yield from module.named_parameters()
-    if envs.FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS.get():
-        for name, buf in module.named_buffers():
-            if buf is not None and buf.numel() * buf.element_size() >= _BUFFER_OFFLOAD_MIN_BYTES:
-                yield name, buf
+    offload_buffers = envs.FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS.get()
+    for name, buf in module.named_buffers():
+        # Empty placement markers follow weights, even when other buffers stay resident.
+        is_weight_device = name.rsplit(".", 1)[-1] == "weight_device" and buf.numel() == 0
+        if is_weight_device or (offload_buffers and buf.numel() * buf.element_size() >= _BUFFER_OFFLOAD_MIN_BYTES):
+            yield name, buf
 
 
 class LayerwiseOffloadState:

fastvideo/layers/lora/linear.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
 
+from dataclasses import dataclass
 import math
 
 import torch
@@ -65,6 +66,15 @@ def requires_side_branch(layer: nn.Module) -> bool:
     return getattr(layer, "weight", None) is None
 
 
+@dataclass(frozen=True, eq=False)
+class LoRAContribution:
+    """A scaled low-rank update for one layer."""
+
+    lora_A: torch.Tensor
+    lora_B: torch.Tensor
+    scale: float
+
+
 class BaseLayerWithLoRA(nn.Module):
     # Wrappers whose forward does not add the unmerged delta must keep merging.
     supports_side_branch: bool = True
@@ -84,7 +94,12 @@ class BaseLayerWithLoRA(nn.Module):
         # A frozen CPU snapshot is needed only to undo inference-time merging.
         # Training never merges adapters; retaining every H3 projection here
         # otherwise adds several GiB of host memory per HSDP rank.
-        self.cpu_weight = (None if training_mode or self.side_branch else base_layer.weight.detach().to("cpu").clone())
+        if training_mode or self.side_branch:
+            self.cpu_weight = None
+        else:
+            self.cpu_weight = base_layer.weight.detach().to("cpu", copy=True)
+        # None disables history; [] starts an eligible layer.
+        self._merged_loras: list[LoRAContribution] | None = [] if self._should_save_merged_loras() else None
         # indicates adapter weights don't contain this layer
         # (which shouldn't normally happen, but we want to separate it from the case of erroneous merging)
         self.disable_lora: bool = False
@@ -122,6 +137,16 @@ class BaseLayerWithLoRA(nn.Module):
             self.lora_A = None
             self.lora_B = None
 
+    def _should_save_merged_loras(self) -> bool:
+        if self.cpu_weight is None or isinstance(self.cpu_weight, DTensor):
+            return False
+        from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
+
+        method = getattr(self.base_layer, "quant_method", None)
+        return (isinstance(method, NVFP4QuantizeMethod)
+                and getattr(method, "_retain_original_weights", None) is not True
+                and not getattr(method, "_is_refine_only_layer", False))
+
     @property
     def lora_strength(self) -> float:
         return self._lora_strength
@@ -207,6 +232,23 @@ class BaseLayerWithLoRA(nn.Module):
                          lora_path: str | None = None,
                          strength: float = 1.0,
                          accumulate: bool = False) -> None:
+        if not math.isfinite(strength) or (lora_alpha is not None and not math.isfinite(lora_alpha)):
+            raise ValueError("LoRA alpha and strength must be finite")
+        if (isinstance(self.base_layer, LinearBase | nn.Linear)
+                and not isinstance(self.base_layer, MergedColumnParallelLinear | QKVParallelLinear)):
+            if A.ndim != 2 or B.ndim != 2 or A.shape[0] == 0 or B.shape[1] != A.shape[0]:
+                raise ValueError("LoRA factors must be matrices with matching nonzero rank")
+            if isinstance(self.base_layer, nn.Linear):
+                input_size, output_size = self.base_layer.in_features, self.base_layer.out_features
+            else:
+                input_size, output_size = self.base_layer.input_size, self.base_layer.output_size
+            if A.shape[1] != input_size or B.shape[0] != output_size:
+                raise ValueError("LoRA factor dimensions do not match the base weight")
+        if not training_mode and not self.side_branch and self.base_layer.weight is None:
+            if accumulate and self.merged:
+                self._restore_merged_weight()
+            else:
+                self._reset_to_base_weight()
         if self.side_branch:
             if accumulate and self.lora_A is not None:
                 raise ValueError("Accumulating adapters is unsupported on a side-branch LoRA layer (packed or "
@@ -236,9 +278,20 @@ class BaseLayerWithLoRA(nn.Module):
         if self.disable_lora or self.side_branch:
             return
 
+        assert self.lora_A is not None and self.lora_B is not None, "LoRA weights not set. Please set them first."
         if self.merged and not accumulate:
             self.unmerge_lora_weights()
-        assert self.lora_A is not None and self.lora_B is not None, "LoRA weights not set. Please set them first."
+        elif self.base_layer.weight is None:
+            if self.merged:
+                self._restore_merged_weight()
+            else:
+                self._reset_to_base_weight()
+        contribution = None
+        if self._merged_loras is not None:
+            scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
+            contribution = LoRAContribution(lora_A=self.lora_A.detach().to("cpu", copy=True),
+                                            lora_B=self.lora_B.detach().to("cpu", copy=True),
+                                            scale=scale * self.lora_strength)
         if isinstance(self.base_layer.weight, DTensor):
             mesh = self.base_layer.weight.data.device_mesh
             unsharded_base_layer = ReplicatedLinear(
@@ -274,18 +327,43 @@ class BaseLayerWithLoRA(nn.Module):
             self.base_layer.weight.data = self._merged_weight(data).to(current_device, non_blocking=True)
 
         self.merged = True
+        if contribution is not None:
+            self._merged_loras.append(contribution)
 
-    def _merged_weight(self, weight: torch.Tensor) -> torch.Tensor:
+    def _merged_weight(self, weight: torch.Tensor, contribution: LoRAContribution | None = None) -> torch.Tensor:
         """``weight`` plus the scaled adapter delta, in ``weight``'s dtype.
 
         An adapter stored wider than the base (an FP32 adapter on a BF16 DiT) is multiplied and added at its own
         precision and rounded once, as a merge-then-cast reference does; with FSDP or without.
         """
-        accumulate_dtype = torch.promote_types(weight.dtype, self.lora_A.dtype)
-        lora_B = self.slice_lora_b_weights(self.lora_B.to(device=weight.device, dtype=accumulate_dtype))
-        lora_A = self.slice_lora_a_weights(self.lora_A.to(device=weight.device, dtype=accumulate_dtype))
-        scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
-        return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * (scale * self.lora_strength)).to(weight.dtype)
+        if contribution is None:
+            scale = ((self.lora_alpha / self.lora_rank) if (self.lora_alpha and self.lora_rank) else 1.0)
+            adapter_a, adapter_b, scale = self.lora_A, self.lora_B, scale * self.lora_strength
+        else:
+            adapter_a, adapter_b, scale = contribution.lora_A, contribution.lora_B, contribution.scale
+        accumulate_dtype = torch.promote_types(weight.dtype, adapter_a.dtype)
+        lora_B = self.slice_lora_b_weights(adapter_b.to(device=weight.device, dtype=accumulate_dtype))
+        lora_A = self.slice_lora_a_weights(adapter_a.to(device=weight.device, dtype=accumulate_dtype))
+        return (weight.to(accumulate_dtype) + (lora_B @ lora_A) * scale).to(weight.dtype)
+
+    @torch.no_grad()
+    def _restore_merged_weight(self) -> None:
+        """Rebuild the current merged weight from immutable CPU contributions."""
+        contributions = self._merged_loras
+        if not contributions:
+            raise RuntimeError("Restoring accumulated LoRA requires saved CPU contributions")
+        previous_weight = self.base_layer.weight
+        try:
+            self._restore_base_weight()
+            current_device = self.base_layer.weight.device
+            data = self.base_layer.weight.data.to(get_local_torch_device())
+            for contribution in contributions:
+                data = self._merged_weight(data, contribution)
+            self.base_layer.weight.data = data.to(current_device)
+        except Exception:
+            self.base_layer.weight = previous_weight
+            raise
+        self.merged = True
 
     @torch.no_grad()
     # @torch.compile(dynamic=True)
@@ -298,18 +376,39 @@ class BaseLayerWithLoRA(nn.Module):
                 return
             raise ValueError("unmerge_lora_weights called but no LoRA is currently merged")
 
-        # avoid precision loss
+        self._reset_to_base_weight()
+
+    @torch.no_grad()
+    def _reset_to_base_weight(self) -> None:
+        self._restore_base_weight()
+        self.merged = False
+        if self._merged_loras is not None:
+            self._merged_loras = []
+
+    @torch.no_grad()
+    def _restore_base_weight(self) -> None:
+        """Restore the base tensor from the immutable CPU snapshot.
+
+        The pipeline calls this inside the block's offload mutation scope; the placement buffer follows its weights.
+        """
         if self.cpu_weight is None:
             raise RuntimeError("Training-mode LoRA layers do not retain an inference unmerge snapshot")
-        if isinstance(self.base_layer.weight, DTensor):
+        if self.base_layer.weight is None:
+            if isinstance(self.cpu_weight, DTensor):
+                raise RuntimeError("Restoring a removed LoRA weight requires an unsharded CPU base snapshot.")
+            weight_device = getattr(self.base_layer, "weight_device", None)
+            if weight_device is None:
+                raise RuntimeError("Restoring a removed LoRA weight requires its weight_device placement buffer.")
+            restored_weight = self.cpu_weight.to(weight_device.device, non_blocking=True, copy=True)
+            self.base_layer.weight = nn.Parameter(restored_weight, requires_grad=False)
+        elif isinstance(self.base_layer.weight, DTensor):
             device = self.base_layer.weight.data.device
-            self.base_layer.weight = nn.Parameter(self.cpu_weight.to(device, non_blocking=True, copy=True))
+            self.base_layer.weight = nn.Parameter(self.cpu_weight.to(device, non_blocking=True, copy=True),
+                                                  requires_grad=self.base_layer.weight.requires_grad)
         else:
             # copy=True: on the snapshot's own device .to() would alias it, and the next merge would write into it.
             self.base_layer.weight.data = self.cpu_weight.data.to(self.base_layer.weight, non_blocking=True, copy=True)
 
-        self.merged = False
-
 
 class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
     """

fastvideo/layers/quantization/nvfp4_config.py

@@ -884,6 +884,10 @@ def convert_model_to_nvfp4(model: torch.nn.Module) -> None:
             else:
                 purged_bytes += weight.numel() * weight.element_size()
                 purged += 1
+                # Preserve weight placement after purge.
+                mod.register_buffer("weight_device",
+                                    torch.empty(0, device=weight.device, dtype=weight.dtype),
+                                    persistent=False)
                 mod.register_parameter("weight", None)
 
     if purged or retained:

fastvideo/pipelines/lora_pipeline.py

@@ -32,49 +32,6 @@ from fastvideo.utils import maybe_download_lora
 logger = init_logger(__name__)
 
 
-def _merge_mode_modules(pipeline: Any, transformer_modules: dict[str, nn.Module]) -> Iterable[nn.Module]:
-    """Yield the modules an adapter change has to rewrite in place.
-
-    Side-branch LoRA layers leave their base weights alone, so packed NVFP4/FP8 linears behind them never
-    block an adapter or strength change. In a transformer that has LoRA layers only the merge-mode ones can
-    block it; in a transformer without them, every module is a candidate.
-    """
-    lora_layers = getattr(pipeline, "lora_layers", None) or {}
-    for transformer_name, transformer in transformer_modules.items():
-        if is_lazy_module(transformer) and not transformer.is_materialized:
-            # Nothing resident to rewrite, and touching the proxy would load it, or reload it from inside its own
-            # materialize transform. A reload re-applies the adapter to fresh layers.
-            continue
-        transformer_layers = lora_layers.get(transformer_name)
-        if transformer_layers is None:
-            yield from transformer.modules()
-            continue
-        for _, layers in transformer_layers.lora_layers_by_block():
-            for layer in layers.values():
-                if not getattr(layer, "side_branch", False):
-                    yield getattr(layer, "base_layer", layer)
-
-
-def _has_nvfp4_weights_without_bf16(pipeline: Any, transformer_modules: dict[str, nn.Module]) -> bool:
-    """Return whether NVFP4 quantization removed the BF16 weight of a layer an adapter must merge into."""
-    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
-
-    return any(
-        isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
-        and getattr(module, "_nvfp4_weight", None) is not None and getattr(module, "weight", None) is None
-        for module in _merge_mode_modules(pipeline, transformer_modules))
-
-
-def _has_quantized_nvfp4_weights(pipeline: Any, transformer_modules: dict[str, nn.Module]) -> bool:
-    """Return whether a layer an adapter must merge into contains an initialized NVFP4 weight."""
-    from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
-
-    return any(
-        isinstance(getattr(module, "quant_method", None), NVFP4QuantizeMethod)
-        and getattr(module, "_nvfp4_weight", None) is not None
-        for module in _merge_mode_modules(pipeline, transformer_modules))
-
-
 def _get_hook_ctx(module: nn.Module | None):
     if module is None:
         return nullcontext()
@@ -626,16 +583,8 @@ class LoRAPipeline(ComposedPipelineBase):
             return
 
         if not self.lora_initialized:
-            # Wrap the layers first, so the NVFP4 check skips packed linears that will run as side branches.
             self.convert_to_lora_layers()
 
-        if not self._setting_constructor_adapter and _has_nvfp4_weights_without_bf16(
-                self, self.trainable_transformer_modules):
-            # TODO(David): Restore the BF16 weights and requantize them after an NVFP4 LoRA adapter change.
-            raise RuntimeError(
-                "Runtime LoRA adapter changes are unsupported after NVFP4 quantization removed the BF16 weights. "
-                "Create a new VideoGenerator with the desired LoRA adapter.")
-
         if not self._setting_constructor_adapter:
             if self._constructor_dense_lora_path is not None:
                 raise RuntimeError(
@@ -806,16 +755,7 @@ class LoRAPipeline(ComposedPipelineBase):
             self._runtime_adapter = (lora_nickname, strength)
 
     def merge_lora_weights(self) -> None:
-        """Merge LoRA weights, then repack the packed quantized buffers for the merged state.
-
-        NVFP4 layers are only supported while their BF16 weight is retained: a purged
-        BF16 weight cannot be unmerged and remerged.
-        """
-        if _has_nvfp4_weights_without_bf16(self, self.trainable_transformer_modules):
-            # TODO(David): Merge from the packed FP4 weight instead of rejecting.
-            raise RuntimeError(
-                "LoRA merge is unsupported after NVFP4 weight quantization removed the BF16 weights because the "
-                "FP4 weights would not reflect the merged LoRA adapter.")
+        """Merge LoRA weights, then repack the packed quantized buffers for the merged state."""
         for (
                 transformer_name,
                 transformer_lora_layers,
@@ -857,15 +797,7 @@ class LoRAPipeline(ComposedPipelineBase):
         self._apply_constructor_adapter()
 
     def unmerge_lora_weights(self) -> None:
-        """Unmerge LoRA weights, then repack the packed quantized buffers for the unmerged state.
-
-        NVFP4 is rejected: its FP4 buffers cannot be rebuilt from a purged BF16 weight.
-        """
-        if _has_quantized_nvfp4_weights(self, self.trainable_transformer_modules):
-            # TODO(David): Preserve BF16 weights and requantize NVFP4 weights after LoRA unmerge.
-            raise RuntimeError(
-                "LoRA unmerge is unsupported after NVFP4 weight quantization because the quantized weights would "
-                "not reflect the unmerged LoRA adapter.")
+        """Unmerge LoRA weights, then repack the packed quantized buffers for the unmerged state."""
         for (
                 transformer_name,
                 transformer_lora_layers,

fastvideo/tests/hooks/test_layerwise_offload.py

@@ -6,11 +6,47 @@ import torch.nn as nn
 
 from fastvideo.hooks.layerwise_offload import (
     LayerwiseOffloadHook,
+    LayerwiseOffloadState,
     enable_layerwise_offload,
 )
 from fastvideo.hooks.hooks import ModuleHookManager
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Placement tracking needs a real CUDA offload stream")
+@pytest.mark.parametrize("offload_buffers", [False, True])
+def test_empty_weight_device_tracks_offload_mutation_and_module_moves(env_overrides, offload_buffers: bool):
+    import fastvideo.envs as envs
+
+    env_overrides.enter_context(envs.FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS.override(offload_buffers))
+    device = torch.device("cuda", torch.cuda.current_device())
+    module = nn.Linear(16, 16, dtype=torch.bfloat16)
+    module.register_buffer("weight_device", torch.empty(0, dtype=torch.bfloat16), persistent=False)
+    module.register_buffer("packed", torch.ones(1 << 20, dtype=torch.uint8))
+    state = LayerwiseOffloadState(torch.cuda.Stream(), device)
+    hook = LayerwiseOffloadHook(state)
+    manager = ModuleHookManager.get_from_or_default(module)
+    manager.append_forward_hook(hook)
+    try:
+        assert module.weight_device.device == module.weight.device == device
+        assert "weight_device" in state.cpu_named_parameters
+        assert ("packed" in state.cpu_named_parameters) is offload_buffers
+        module.to(device)
+        assert module.weight_device.device == device
+        module.to("cpu")
+        assert module.weight_device.device.type == "cpu"
+        with hook.mutate_params_scope():
+            assert module.weight_device.device == module.weight.device == device
+        state.prefetch_params()
+        state.wait_and_replace_params()
+        assert module.weight_device.device == module.weight.device == device
+        state.release_gpu_params()
+        assert module.weight_device.device == module.weight.device == device
+        assert module.weight_device.untyped_storage().nbytes() == 0
+        assert "weight_device" not in module.state_dict()
+    finally:
+        manager.remove_forward_hook(hook.name())
+
+
 class SimpleBlock(nn.Module):
     """A simple block with linear layers for testing."""
 

fastvideo/tests/layers/test_lora_side_branch.py

@@ -331,6 +331,67 @@ def test_side_branch_rejects_accumulating_a_second_adapter() -> None:
         layer.set_lora_weights(a, b, accumulate=True)
 
 
+@pytest.mark.parametrize("mode", ["dense", "training", "side_branch"])
+def test_invalid_replacement_keeps_the_current_adapter(mode: str) -> None:
+    training = mode == "training"
+    base = _packed_linear() if mode == "side_branch" else _dense_linear()
+    layer = get_lora_layer(base, lora_rank=RANK, lora_alpha=RANK, training_mode=training)
+    a, b = _factors()
+    layer.set_lora_weights(a, b, lora_alpha=RANK, strength=0.5, training_mode=training)
+    previous_a, previous_b = layer.lora_A, layer.lora_B
+    inputs = _x()
+    output = layer(inputs)[0].detach().clone()
+    assert (layer.cpu_weight is None) == (training or mode == "side_branch")
+
+    invalid = [
+        (a, b, float("nan"), 1.0),
+        (a, b, RANK, float("inf")),
+        (a.unsqueeze(0), b, RANK, 1.0),
+        (a[:0], b[:, :0], RANK, 1.0),
+        (a, b[:, :-1], RANK, 1.0),
+        (a[:, :-1], b, RANK, 1.0),
+        (a, b[:-1], RANK, 1.0),
+    ]
+    for new_a, new_b, alpha, strength in invalid:
+        with pytest.raises(ValueError):
+            layer.set_lora_weights(new_a, new_b, lora_alpha=alpha, strength=strength, training_mode=training)
+        assert layer.lora_A is previous_a and layer.lora_B is previous_b
+        assert layer.lora_strength == 0.5 and layer.lora_alpha == RANK
+        torch.testing.assert_close(layer(inputs)[0], output, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("axis", ["row", "column"])
+def test_parallel_layer_merges_full_factors_into_local_partition(monkeypatch, axis: str) -> None:
+    from fastvideo.layers import linear
+
+    monkeypatch.setattr(linear, "get_tp_world_size", lambda: 2)
+    monkeypatch.setattr(linear, "get_tp_rank", lambda: 1)
+    monkeypatch.setattr(lora_linear, "get_tp_rank", lambda: 1)
+    projection = linear.RowParallelLinear if axis == "row" else linear.ColumnParallelLinear
+    base = projection(IN, OUT, bias=False, params_dtype=torch.float32)
+    base.weight.data.zero_()
+    layer = get_lora_layer(base)
+    a, b = _factors()
+    layer.set_lora_weights(a, b)
+    full_delta = b @ a
+    expected = full_delta[:, IN // 2:] if axis == "row" else full_delta[OUT // 2:, :]
+    torch.testing.assert_close(layer.base_layer.weight, expected)
+
+
+def test_fused_training_layer_accepts_stacked_projection_factors(monkeypatch) -> None:
+    from fastvideo.layers import linear
+
+    monkeypatch.setattr(linear, "get_tp_world_size", lambda: 1)
+    base = linear.MergedColumnParallelLinear(IN, [OUT, OUT], bias=False, params_dtype=torch.bfloat16)
+    layer = get_lora_layer(base, lora_rank=RANK, lora_alpha=RANK, training_mode=True)
+    a, b = _factors()
+    fused_b = torch.stack([b, b * 2])
+    layer.set_lora_weights(a, fused_b, training_mode=True)
+    torch.testing.assert_close(layer.lora_B, fused_b)
+    with pytest.raises(ValueError, match="finite"):
+        layer.set_lora_weights(a, fused_b, strength=float("nan"), training_mode=True)
+
+
 def test_dense_linear_still_merges_and_unmerges_exactly() -> None:
     base = _dense_linear()
     original = base.weight.detach().clone()
@@ -351,28 +412,16 @@ def test_parallel_wrappers_keep_merging() -> None:
     assert not RowParallelLinearWithLoRA.supports_side_branch
 
 
-def _pipeline_over(layer: BaseLayerWithLoRA) -> "SimpleNamespace":
-    from types import SimpleNamespace
-    transformer = torch.nn.Module()
-    transformer.layer = layer
-    layers = SimpleNamespace(lora_layers_by_block=lambda: iter([(None, {"layer": layer})]))
-    return SimpleNamespace(lora_layers={"transformer": layers}, trainable_transformer_modules={"transformer": transformer})
-
-
-def test_packed_nvfp4_layers_behind_a_side_branch_do_not_block_adapter_changes() -> None:
+def test_packed_nvfp4_side_branch_can_skip_repacking() -> None:
     from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
-    from fastvideo.pipelines.lora_pipeline import _all_side_branch, _has_nvfp4_weights_without_bf16
+    from fastvideo.pipelines.lora_pipeline import _all_side_branch
 
     base = _packed_linear()
     base.quant_method = object.__new__(NVFP4QuantizeMethod)
-    pipeline = _pipeline_over(get_lora_layer(base))
-
-    assert not _has_nvfp4_weights_without_bf16(pipeline, pipeline.trainable_transformer_modules)
-    assert _all_side_branch([pipeline.lora_layers["transformer"].lora_layers_by_block().__next__()[1]["layer"]])
+    assert _all_side_branch([get_lora_layer(base)])
 
 
 def test_runtime_adapter_on_packed_nvfp4_without_constructor_lora_is_not_blocked() -> None:
-    """Wrapping runs before the NVFP4 check, so packed linears that become side branches do not block it."""
     from types import SimpleNamespace
 
     from fastvideo.layers.quantization.nvfp4_config import NVFP4QuantizeMethod
@@ -400,7 +449,6 @@ def test_runtime_adapter_on_packed_nvfp4_without_constructor_lora_is_not_blocked
         pipeline.lora_initialized = True
 
     pipeline.convert_to_lora_layers = convert_to_lora_layers
-    # The adapter is unknown and has no path, so the call stops at the lookup that follows the NVFP4 check.
     with pytest.raises(ValueError, match="not found in the pipeline"):
         LoRAPipeline.set_lora_adapter(pipeline, "runtime-adapter")
     assert transformer.layer.side_branch

fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py

@@ -47,8 +47,9 @@ def test_dense_transformer_safetensors_drops_packed_export(tmp_path) -> None:
     assert nv.find_minimax_h3_nvfp4_dit_export([str(shard)]) == str(packed)
 
 
-def test_load_minimax_h3_nvfp4_dit_export_overlays_and_purges_weight(tmp_path) -> None:
-    model = _model()
+@pytest.mark.parametrize("dense_device", ["cpu", "meta"])
+def test_load_minimax_h3_nvfp4_dit_export_overlays_and_purges_weight(tmp_path, dense_device: str) -> None:
+    model = _model().to(dense_device)
     linear = model.transformer_blocks[0].attn.to_q
     prefix = "transformer_blocks.0.attn.to_q"
     tensors = {
@@ -64,6 +65,7 @@ def test_load_minimax_h3_nvfp4_dit_export_overlays_and_purges_weight(tmp_path) -
 
     assert loaded == 1
     assert linear.weight is None
+    assert linear._nvfp4_weight.device.type == "cpu"
     assert linear._nvfp4_weight.dtype is torch.uint8
     assert tuple(linear._nvfp4_weight.shape) == (8, 8)
     assert linear._nvfp4_alpha.item() == 0.5

fastvideo/tests/ops/quantization/test_nvfp4_lora_accumulation.py

@@ -0,0 +1,284 @@
+# SPDX-License-Identifier: Apache-2.0
+"""CPU contracts for replaying active LoRA contributions after an NVFP4 purge."""
+
+from __future__ import annotations
+
+from dataclasses import FrozenInstanceError
+
+import pytest
+import torch
+
+from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.lora.linear import get_lora_layer
+from fastvideo.layers.quantization.mxfp8_config import MXFP8QuantizeMethod
+from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+from fastvideo.tests.ops.quantization.test_nvfp4_minimax_h3_wiring import (
+    _assert_cpu_packed_weight,
+    _CpuPackedNVFP4Method,
+    _cpu_pack_nvfp4,
+    _nvfp4_lifecycle_pipeline,
+    cpu_nvfp4_lifecycle,
+)
+
+
+def _reference_update(weight, factors, *, alpha=None, strength=1.0):
+    a, b = factors
+    rank = a.shape[0]
+    scale = ((alpha / rank) if alpha is not None else 1.0) * strength
+    dtype = torch.promote_types(weight.dtype, a.dtype)
+    delta = b.to(dtype) @ a.to(dtype)
+    return (weight.to(dtype) + delta * scale).to(weight.dtype)
+
+
+def _install(pipeline, name, factors, *, alpha=None, strength=1.0, accumulate=True):
+    a, b = factors
+    pipeline.lora_adapters[name] = {"linear.lora_A": a, "linear.lora_B": b}
+    if alpha is not None:
+        pipeline.lora_adapters[name]["linear.lora_alpha"] = alpha
+    LoRAPipeline.set_lora_adapter(pipeline, name, strength=strength, accumulate=accumulate)
+
+
+def test_accumulation_replays_multiple_adapters_with_original_dtype_and_scale(cpu_nvfp4_lifecycle):
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(False)
+    third = (torch.randn(2, 32) / 8, torch.randn(32, 2) / 4)
+    expected = _reference_update(original, first)
+
+    for name, factors, alpha, strength in [
+        ("second", second, 4.0, 0.5),
+        ("third", third, 3.5, -0.75),
+        ("third", third, 3.5, -0.75),
+    ]:
+        expected = _reference_update(expected, factors, alpha=alpha, strength=strength)
+        _install(pipeline, name, factors, alpha=alpha, strength=strength)
+        _assert_cpu_packed_weight(wrapped, expected, False)
+
+    records = wrapped._merged_loras
+    assert len(records) == 4
+    assert records[0].lora_A.dtype == torch.bfloat16
+    assert records[2].lora_A.dtype == records[2].lora_B.dtype == torch.float32
+    assert records[2].scale == (3.5 / 2) * -0.75
+    assert all(record.lora_A.device.type == record.lora_B.device.type == "cpu" for record in records)
+    with pytest.raises(FrozenInstanceError):
+        records[0].scale = 0.0
+    torch.testing.assert_close(wrapped.cpu_weight, original, rtol=0, atol=0)
+
+    _install(pipeline, "replacement", third, alpha=3.5, strength=-0.75, accumulate=False)
+    _assert_cpu_packed_weight(wrapped, _reference_update(original, third, alpha=3.5, strength=-0.75), False)
+    assert len(wrapped._merged_loras) == 1
+
+
+def test_direct_accumulating_merge_after_each_purge_keeps_previous_updates(cpu_nvfp4_lifecycle):
+    _, wrapped, original, first, _ = _nvfp4_lifecycle_pipeline(False)
+    expected = _reference_update(original, first)
+
+    for _ in range(2):
+        wrapped.merge_lora_weights(accumulate=True)
+        expected = _reference_update(expected, first)
+        torch.testing.assert_close(wrapped.base_layer.weight, expected, rtol=0, atol=0)
+        _cpu_pack_nvfp4(wrapped)
+        _assert_cpu_packed_weight(wrapped, expected, False)
+
+    assert len(wrapped._merged_loras) == 3
+
+
+def test_merge_after_clear_preserves_purged_state(cpu_nvfp4_lifecycle):
+    pipeline, wrapped, original, _, _ = _nvfp4_lifecycle_pipeline(False)
+    LoRAPipeline.clear_lora_adapter(pipeline)
+    assert wrapped.base_layer.weight is None
+    previous_flags = (wrapped.merged, wrapped.disable_lora, wrapped.side_branch)
+    previous_records = wrapped._merged_loras
+    buffers = {name: (value, value.clone()) for name, value in wrapped.base_layer.named_buffers()}
+
+    with pytest.raises(AssertionError, match="LoRA weights not set"):
+        LoRAPipeline.merge_lora_weights(pipeline)
+
+    assert wrapped.base_layer.weight is None
+    assert (wrapped.merged, wrapped.disable_lora, wrapped.side_branch) == previous_flags
+    assert wrapped.lora_A is None and wrapped.lora_B is None
+    assert wrapped._merged_loras is previous_records
+    torch.testing.assert_close(wrapped.cpu_weight, original, rtol=0, atol=0)
+    for name, (previous_buffer, expected) in buffers.items():
+        assert getattr(wrapped.base_layer, name) is previous_buffer
+        torch.testing.assert_close(getattr(wrapped.base_layer, name), expected, rtol=0, atol=0)
+
+
+def test_accumulated_round_trip_repacking_matches_dense_lifecycle(cpu_nvfp4_lifecycle):
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(False)
+    base = ReplicatedLinear(32, 32, bias=False, params_dtype=torch.bfloat16)
+    base.weight.data.copy_(original)
+    dense = get_lora_layer(base)
+    dense.set_lora_weights(*(factor.clone() for factor in first), lora_alpha=4)
+    dense.set_lora_weights(*(factor.clone() for factor in second), lora_alpha=4, strength=0.5, accumulate=True)
+    _install(pipeline, "second", second, alpha=4, strength=0.5)
+    _assert_cpu_packed_weight(wrapped, dense.base_layer.weight, False)
+
+    for _ in range(2):
+        dense.unmerge_lora_weights()
+        LoRAPipeline.unmerge_lora_weights(pipeline)
+        assert wrapped.merged == dense.merged
+        _assert_cpu_packed_weight(wrapped, dense.base_layer.weight, False)
+
+        dense.merge_lora_weights()
+        LoRAPipeline.merge_lora_weights(pipeline)
+        assert wrapped.merged == dense.merged
+        _assert_cpu_packed_weight(wrapped, dense.base_layer.weight, False)
+
+
+def test_replay_rounds_to_bf16_after_each_contribution(cpu_nvfp4_lifecycle):
+    base = ReplicatedLinear(32, 32, bias=False, params_dtype=torch.bfloat16)
+    base.weight.data.fill_(256)
+    base.quant_method = _CpuPackedNVFP4Method(False)
+    wrapped = get_lora_layer(base)
+    factors = (torch.ones(1, 32, dtype=torch.bfloat16), torch.ones(32, 1, dtype=torch.bfloat16))
+
+    wrapped.set_lora_weights(*factors, lora_alpha=1)
+    _cpu_pack_nvfp4(wrapped)
+    wrapped.set_lora_weights(*factors, lora_alpha=1, accumulate=True)
+    _cpu_pack_nvfp4(wrapped)
+    wrapped.set_lora_weights(*factors, lora_alpha=1, strength=0.0, accumulate=True)
+
+    # 256 + 1 rounds back to 256 at each BF16 merge; combining the two deltas gives 258.
+    expected = torch.full((32, 32), 256, dtype=torch.bfloat16)
+    torch.testing.assert_close(wrapped.base_layer.weight, expected, rtol=0, atol=0)
+    assert not torch.equal(wrapped.base_layer.weight, expected + 2)
+
+
+def test_cpu_contributions_do_not_alias_adapter_cache_or_follow_scalar_changes(cpu_nvfp4_lifecycle):
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(False)
+    applied_first = tuple(t.clone() for t in first)
+    records = wrapped._merged_loras
+    saved_a, saved_b, saved_scale = records[0].lora_A, records[0].lora_B, records[0].scale
+    assert saved_a.data_ptr() != first[0].data_ptr()
+    assert saved_b.data_ptr() != first[1].data_ptr()
+
+    first[0].fill_(9)
+    first[1].fill_(-7)
+    wrapped.lora_strength = 11.0
+    wrapped.lora_alpha = 13.0
+    _install(pipeline, "second", second, alpha=4, strength=0.5)
+
+    expected = _reference_update(_reference_update(original, applied_first), second, alpha=4, strength=0.5)
+    _assert_cpu_packed_weight(wrapped, expected, False)
+    torch.testing.assert_close(saved_a, applied_first[0], rtol=0, atol=0)
+    torch.testing.assert_close(saved_b, applied_first[1], rtol=0, atol=0)
+    assert saved_scale == 1.0
+
+
+@pytest.mark.parametrize("retain", [False, True])
+def test_tensor_restore_preserves_state_until_explicit_reset(cpu_nvfp4_lifecycle, retain: bool):
+    _, wrapped, original, _, _ = _nvfp4_lifecycle_pipeline(retain)
+    previous_records = wrapped._merged_loras
+    previous_a, previous_b = wrapped.lora_A, wrapped.lora_B
+
+    wrapped._restore_base_weight()
+
+    torch.testing.assert_close(wrapped.base_layer.weight, original, rtol=0, atol=0)
+    assert wrapped.merged
+    assert wrapped._merged_loras is previous_records
+    assert wrapped.lora_A is previous_a and wrapped.lora_B is previous_b
+
+    wrapped._reset_to_base_weight()
+
+    assert not wrapped.merged
+    if retain:
+        assert wrapped._merged_loras is None
+    else:
+        assert wrapped._merged_loras == [] and wrapped._merged_loras is not previous_records
+    torch.testing.assert_close(wrapped.base_layer.weight, original, rtol=0, atol=0)
+    torch.testing.assert_close(wrapped.cpu_weight, original, rtol=0, atol=0)
+    assert wrapped.lora_A is previous_a and wrapped.lora_B is previous_b
+
+
+@pytest.mark.parametrize("invalid", ["rank", "width", "alpha", "strength"])
+@pytest.mark.parametrize("retain", [False, True])
+def test_invalid_accumulation_preserves_previous_factors_history_and_packed_weights(cpu_nvfp4_lifecycle, invalid,
+                                                                                   retain):
+    pipeline, wrapped, _, _, second = _nvfp4_lifecycle_pipeline(retain)
+    _install(pipeline, "second", second, strength=0.5)
+    previous_a, previous_b = wrapped.lora_A, wrapped.lora_B
+    previous_records = wrapped._merged_loras
+    previous_weight = wrapped.base_layer.weight.detach().clone() if retain else None
+    packed = {name: value.clone() for name, value in wrapped.base_layer.named_buffers()}
+    a, b = torch.randn(3, 32), torch.randn(32, 3)
+    alpha, strength = 3.0, 1.0
+    if invalid == "rank":
+        b = torch.randn(32, 2)
+    elif invalid == "width":
+        a = torch.randn(3, 31)
+    elif invalid == "alpha":
+        alpha = float("nan")
+    else:
+        strength = float("inf")
+
+    with pytest.raises(ValueError):
+        wrapped.set_lora_weights(a, b, lora_alpha=alpha, strength=strength, accumulate=True)
+
+    assert wrapped.merged
+    if retain:
+        torch.testing.assert_close(wrapped.base_layer.weight, previous_weight, rtol=0, atol=0)
+    else:
+        assert wrapped.base_layer.weight is None and len(previous_records) == 2
+    assert wrapped.lora_A is previous_a and wrapped.lora_B is previous_b
+    assert wrapped._merged_loras is previous_records
+    for name, value in packed.items():
+        torch.testing.assert_close(getattr(wrapped.base_layer, name), value, rtol=0, atol=0)
+
+
+def test_failed_replay_restores_previous_state_and_can_be_retried(cpu_nvfp4_lifecycle,
+                                                               monkeypatch: pytest.MonkeyPatch):
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(False)
+    _install(pipeline, "second", second, strength=0.5)
+    third = (torch.randn(2, 32) / 8, torch.randn(32, 2) / 4)
+    previous_a, previous_b = wrapped.lora_A, wrapped.lora_B
+    previous_records = wrapped._merged_loras
+    packed = {name: value.clone() for name, value in wrapped.base_layer.named_buffers()}
+    real_merge = wrapped._merged_weight
+    replayed = 0
+
+    def fail_on_second_contribution(weight, contribution=None):
+        nonlocal replayed
+        if contribution is not None:
+            replayed += 1
+            if replayed == 2:
+                raise RuntimeError("injected replay failure")
+        return real_merge(weight, contribution)
+
+    with monkeypatch.context() as patch:
+        patch.setattr(wrapped, "_merged_weight", fail_on_second_contribution)
+        with pytest.raises(RuntimeError, match="injected replay failure"):
+            wrapped.set_lora_weights(*third, lora_alpha=2, strength=-0.25, accumulate=True)
+
+    assert replayed == 2
+    assert wrapped.base_layer.weight is None and wrapped.merged
+    assert wrapped.lora_A is previous_a and wrapped.lora_B is previous_b
+    assert wrapped._merged_loras is previous_records
+    assert len(previous_records) == 2
+    for name, value in packed.items():
+        torch.testing.assert_close(getattr(wrapped.base_layer, name), value, rtol=0, atol=0)
+
+    wrapped.set_lora_weights(*third, lora_alpha=2, strength=-0.25, accumulate=True)
+    expected = _reference_update(_reference_update(original, first), second, strength=0.5)
+    expected = _reference_update(expected, third, alpha=2, strength=-0.25)
+    _cpu_pack_nvfp4(wrapped)
+    _assert_cpu_packed_weight(wrapped, expected, False)
+    assert len(wrapped._merged_loras) == 3
+
+
+@pytest.mark.parametrize("mode", ["bf16", "mxfp8", "nvfp4_retained", "nvfp4_refine", "training", "side_branch"])
+def test_other_paths_do_not_keep_cpu_contribution_snapshots(cpu_nvfp4_lifecycle, mode):
+    base = ReplicatedLinear(32, 32, bias=False, params_dtype=torch.bfloat16)
+    base.weight.data.fill_(0.125)
+    if mode == "mxfp8":
+        base.quant_method = object.__new__(MXFP8QuantizeMethod)
+    elif mode in ("nvfp4_retained", "nvfp4_refine", "training", "side_branch"):
+        base.quant_method = _CpuPackedNVFP4Method(mode == "nvfp4_retained")
+        if mode == "nvfp4_refine":
+            base.quant_method._is_refine_only_layer = True
+        if mode == "side_branch":
+            _cpu_pack_nvfp4(base)
+    training = mode == "training"
+    wrapped = get_lora_layer(base, lora_rank=2, lora_alpha=2, training_mode=training)
+    factors = (torch.ones(2, 32, dtype=torch.bfloat16) / 8,
+               torch.ones(32, 2, dtype=torch.bfloat16) / 8)
+    wrapped.set_lora_weights(*factors, lora_alpha=2, training_mode=training)
+    assert wrapped._merged_loras is None

fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py

@@ -252,8 +252,9 @@ def test_unmerge_then_merge_lora_weights_mxfp8_requantizes_each_step(
 
 
 @pytest.mark.parametrize("retain_bf16_weight", [False, True])
-def test_unmerge_lora_weights_rejects_quantized_nvfp4(retain_bf16_weight: bool) -> None:
-    """Reject NVFP4 LoRA unmerge whether the BF16 weight was removed or retained."""
+@pytest.mark.parametrize("operation", ["merge_lora_weights", "unmerge_lora_weights"])
+def test_nvfp4_without_lora_layers_is_a_lifecycle_noop(retain_bf16_weight: bool, operation: str,
+                                                      monkeypatch: pytest.MonkeyPatch) -> None:
     from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 
     transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
@@ -262,32 +263,17 @@ def test_unmerge_lora_weights_rejects_quantized_nvfp4(retain_bf16_weight: bool)
         nn.Parameter(torch.empty(1), requires_grad=False) if retain_bf16_weight else None,
     )
     transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1))
-    pipeline = SimpleNamespace(trainable_transformer_modules={"transformer": transformer})
+    pipeline = SimpleNamespace(trainable_transformer_modules={"transformer": transformer}, lora_layers={})
+    previous_weight = transformer.linear.weight
+    previous_packed = transformer.linear._nvfp4_weight
 
-    with pytest.raises(RuntimeError, match="LoRA unmerge is unsupported after NVFP4 weight quantization"):
-        LoRAPipeline.unmerge_lora_weights(pipeline)
+    def unexpected_pack(model):
+        pytest.fail("An adapter-free lifecycle operation must not repack the model")
 
-
-def test_set_lora_adapter_rejects_nvfp4_without_bf16_weight() -> None:
-    """Reject runtime adapter changes after NVFP4 removes the BF16 weight."""
-    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
-
-    transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
-    transformer.linear.register_parameter("weight", None)
-    transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1))
-    pipeline = SimpleNamespace(
-        _setting_constructor_adapter=False,
-        cur_adapter_name="first-adapter",
-        cur_adapter_path="/first-adapter",
-        cur_adapter_strength=1.0,
-        lora_adapter_paths={},
-        lora_initialized=True,
-        lora_layers={},
-        trainable_transformer_modules={"transformer": transformer},
-    )
-
-    with pytest.raises(RuntimeError, match="Runtime LoRA adapter changes are unsupported after NVFP4 quantization"):
-        LoRAPipeline.set_lora_adapter(pipeline, "second-adapter", "/second-adapter")
+    monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", unexpected_pack)
+    getattr(LoRAPipeline, operation)(pipeline)
+    assert transformer.linear.weight is previous_weight
+    assert transformer.linear._nvfp4_weight is previous_packed
 
 
 def test_set_lora_adapter_mxfp8_requantizes_after_adapter_merge(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -320,24 +306,455 @@ def test_set_lora_adapter_mxfp8_requantizes_after_adapter_merge(monkeypatch: pyt
     assert events == ["adapter merged", "weights requantized"]
 
 
-@pytest.mark.parametrize("retain_bf16_weight", [False, True])
-def test_merge_lora_weights_rejects_nvfp4_without_bf16(retain_bf16_weight: bool) -> None:
-    """Reject the NVFP4 LoRA merge only once the BF16 weight is gone."""
+class _CpuPackedNVFP4Method(nvfp4.NVFP4QuantizeMethod):
+    """CPU stand-in for packed storage and parameter purge, not an NVFP4 kernel emulator."""
+
+    def __init__(self, retain: bool) -> None:
+        self.retain = retain
+        self._retain_original_weights = retain
+        self._is_refine_only_layer = False
+
+    def apply(self, layer, x, bias=None):
+        weight = layer._nvfp4_weight.to(x.dtype) * layer._nvfp4_weight_scale.to(x.dtype)
+        return torch.nn.functional.linear(x, weight, bias)
+
+
+def _cpu_packed_state(weight: torch.Tensor) -> dict[str, torch.Tensor]:
+    return {
+        "_nvfp4_weight": torch.round(weight.float() * 32).to(torch.int16),
+        "_nvfp4_weight_scale": torch.full((), 1 / 32),
+        "_weight_global_sf": weight.float().abs().amax().to(torch.bfloat16),
+        "_nvfp4_alpha": weight.float().abs().amax().reciprocal(),
+    }
+
+
+def _cpu_pack_nvfp4(model: nn.Module) -> None:
+    for layer in model.modules():
+        if isinstance(getattr(layer, "quant_method", None), _CpuPackedNVFP4Method) and layer.weight is not None:
+            for name, value in _cpu_packed_state(layer.weight.detach()).items():
+                layer.register_buffer(name, value, persistent=False)
+            if not layer.quant_method.retain:
+                layer.register_buffer("weight_device", torch.empty(0, device=layer.weight.device), persistent=False)
+                layer.register_parameter("weight", None)
+
+
+@pytest.fixture
+def cpu_nvfp4_lifecycle(monkeypatch: pytest.MonkeyPatch):
+    from fastvideo.layers.lora import linear as lora_linear
+
+    monkeypatch.setattr(lora_linear, "get_local_torch_device", lambda: torch.device("cpu"))
+    monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", _cpu_pack_nvfp4)
+    monkeypatch.setattr("fastvideo.pipelines.lora_pipeline.dist.get_rank", lambda: 0)
+    with torch._dynamo.config.patch(disable=True):
+        yield
+
+
+def _nvfp4_lifecycle_pipeline(retain: bool):
+    from fastvideo.layers.linear import ReplicatedLinear
+    from fastvideo.layers.lora.linear import get_lora_layer
+
+    torch.manual_seed(7)
+    base = ReplicatedLinear(32, 32, bias=False, params_dtype=torch.bfloat16)
+    with torch.no_grad():
+        base.weight.copy_(torch.randn_like(base.weight) / 8)
+    original = base.weight.detach().clone()
+    base.quant_method = _CpuPackedNVFP4Method(retain)
+    wrapped = get_lora_layer(base)
+    first = (torch.randn(4, 32, dtype=torch.bfloat16) / 8, torch.randn(32, 4, dtype=torch.bfloat16) / 4)
+    second = (torch.randn(4, 32, dtype=torch.bfloat16) / 4, torch.randn(32, 4, dtype=torch.bfloat16) / 4)
+    wrapped.set_lora_weights(*first, lora_alpha=4)
+    _cpu_pack_nvfp4(wrapped)
+    transformer = nn.Module()
+    transformer.linear = wrapped
+    layer_group = SimpleNamespace(lora_layers_by_block=lambda: iter([(None, {"linear": wrapped})]))
+    pipeline = SimpleNamespace(
+        lora_layers={"transformer": layer_group},
+        trainable_transformer_modules={"transformer": transformer},
+        _setting_constructor_adapter=False,
+        _constructor_dense_lora_path=None,
+        _apply_constructor_adapter=lambda: None,
+        cur_adapter_name="first",
+        cur_adapter_path=None,
+        cur_adapter_strength=1.0,
+        lora_initialized=True,
+        lora_adapter_paths={},
+        lora_adapters={
+            "first": {"linear.lora_A": first[0], "linear.lora_B": first[1]},
+            "second": {"linear.lora_A": second[0], "linear.lora_B": second[1]},
+        },
+        fastvideo_args=SimpleNamespace(training_mode=False),
+    )
+    return pipeline, wrapped, original, first, second
+
+
+def _assert_cpu_packed_weight(wrapped, expected: torch.Tensor, retain: bool) -> None:
+    for name, value in _cpu_packed_state(expected).items():
+        torch.testing.assert_close(getattr(wrapped.base_layer, name), value, rtol=0, atol=0)
+    assert (wrapped.base_layer.weight is not None) is retain
+
+
+@pytest.mark.parametrize("retain", [False, True])
+def test_nvfp4_cpu_backed_unmerge_remerge_rebuilds_every_buffer(cpu_nvfp4_lifecycle, retain: bool) -> None:
     from fastvideo.pipelines.lora_pipeline import LoRAPipeline
 
-    transformer = _TaggedModel(object.__new__(nvfp4.NVFP4QuantizeMethod))
-    transformer.linear.register_parameter(
-        "weight",
-        nn.Parameter(torch.empty(1), requires_grad=False) if retain_bf16_weight else None,
-    )
-    transformer.linear.register_buffer("_nvfp4_weight", torch.empty(1))
-    pipeline = SimpleNamespace(
-        trainable_transformer_modules={"transformer": transformer},
-        lora_layers={},
-    )
-
-    if retain_bf16_weight:
+    pipeline, wrapped, original, first, _ = _nvfp4_lifecycle_pipeline(retain)
+    merged = original + first[1] @ first[0]
+    _assert_cpu_packed_weight(wrapped, merged, retain)
+    inputs = torch.randn(3, 32, dtype=torch.bfloat16)
+    for _ in range(2):
+        LoRAPipeline.unmerge_lora_weights(pipeline)
+        _assert_cpu_packed_weight(wrapped, original, retain)
+        assert not wrapped.merged
+        if retain:
+            assert not wrapped.base_layer.weight.requires_grad
+        packed = _cpu_packed_state(original)
+        expected = torch.nn.functional.linear(
+            inputs, packed["_nvfp4_weight"].to(inputs) * packed["_nvfp4_weight_scale"].to(inputs))
+        expected += inputs @ first[0].T @ first[1].T
+        torch.testing.assert_close(wrapped(inputs)[0], expected, rtol=0, atol=0)
         LoRAPipeline.merge_lora_weights(pipeline)
+        _assert_cpu_packed_weight(wrapped, merged, retain)
+        torch.testing.assert_close(wrapped.cpu_weight, original, rtol=0, atol=0)
+
+
+def test_nvfp4_restore_recreates_a_frozen_inference_parameter(cpu_nvfp4_lifecycle) -> None:
+    _, wrapped, original, _, _ = _nvfp4_lifecycle_pipeline(False)
+    wrapped.unmerge_lora_weights()
+
+    assert not wrapped.base_layer.weight.requires_grad
+    assert wrapped.base_layer.weight.data_ptr() != wrapped.cpu_weight.data_ptr()
+    torch.testing.assert_close(wrapped.base_layer.weight, original, rtol=0, atol=0)
+
+
+def test_lora_training_keeps_adapter_gradients_without_inference_restore(cpu_nvfp4_lifecycle,
+                                                                       monkeypatch: pytest.MonkeyPatch) -> None:
+    """The shared wrapper trains A/B while freezing the base and never restoring inference weights."""
+    from fastvideo.layers.linear import ReplicatedLinear
+    from fastvideo.layers.lora.linear import get_lora_layer
+
+    torch.manual_seed(11)
+    base = ReplicatedLinear(8, 6, bias=False, params_dtype=torch.float32)
+    base.weight.requires_grad_(True)
+    wrapped = get_lora_layer(base, lora_rank=2, lora_alpha=2, training_mode=True)
+    assert wrapped.cpu_weight is None
+
+    def unexpected_restore():
+        pytest.fail("LoRA training must not restore or merge inference weights")
+
+    monkeypatch.setattr(wrapped, "_restore_base_weight", unexpected_restore)
+    wrapped.set_lora_weights(torch.randn(2, 8), torch.randn(6, 2), training_mode=True)
+    assert not wrapped.merged and not wrapped.base_layer.weight.requires_grad
+    assert wrapped.lora_A.requires_grad and wrapped.lora_B.requires_grad
+    original = wrapped.base_layer.weight.detach().clone()
+    adapter_before = [wrapped.lora_A.detach().clone(), wrapped.lora_B.detach().clone()]
+    optimizer = torch.optim.SGD([wrapped.lora_A, wrapped.lora_B], lr=0.01)
+    inputs = torch.randn(3, 8, requires_grad=True)
+
+    wrapped(inputs)[0].square().mean().backward()
+
+    assert wrapped.base_layer.weight.grad is None
+    assert inputs.grad is not None and inputs.grad.abs().sum() > 0
+    for parameter in (wrapped.lora_A, wrapped.lora_B):
+        assert parameter.grad is not None and parameter.grad.abs().sum() > 0
+    optimizer.step()
+    torch.testing.assert_close(wrapped.base_layer.weight, original, rtol=0, atol=0)
+    for parameter, previous in zip((wrapped.lora_A, wrapped.lora_B), adapter_before, strict=True):
+        assert not torch.equal(parameter, previous)
+
+
+@pytest.mark.parametrize("retain", [False, True])
+def test_nvfp4_cpu_backed_adapter_and_strength_changes_replace_old_delta(cpu_nvfp4_lifecycle, retain: bool) -> None:
+    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(retain)
+    for nickname, strength, factors in [("second", 0.5, second), ("second", 1.5, second), ("first", 1.0, first)]:
+        LoRAPipeline.set_lora_adapter(pipeline, nickname, strength=strength)
+        _assert_cpu_packed_weight(wrapped, original + (factors[1] @ factors[0]) * strength, retain)
+        torch.testing.assert_close(wrapped.cpu_weight, original, rtol=0, atol=0)
+    LoRAPipeline.clear_lora_adapter(pipeline)
+    _assert_cpu_packed_weight(wrapped, original, retain)
+    assert wrapped.lora_A is None and wrapped.lora_B is None
+
+
+def test_nvfp4_retained_bf16_accumulation_requantizes(cpu_nvfp4_lifecycle) -> None:
+    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(True)
+    LoRAPipeline.set_lora_adapter(pipeline, "second", strength=0.5, accumulate=True)
+    _assert_cpu_packed_weight(wrapped, (original + first[1] @ first[0]) + (second[1] @ second[0]) * 0.5, True)
+
+
+@pytest.mark.parametrize("entrypoint", ["pipeline", "layer", "merge"])
+def test_nvfp4_purged_accumulation_reconstructs_before_appending(cpu_nvfp4_lifecycle, entrypoint: str) -> None:
+    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+
+    pipeline, wrapped, original, first, second = _nvfp4_lifecycle_pipeline(False)
+    expected = original + first[1] @ first[0]
+    if entrypoint == "pipeline":
+        LoRAPipeline.set_lora_adapter(pipeline, "second", accumulate=True)
+    elif entrypoint == "layer":
+        wrapped.set_lora_weights(*second, accumulate=True)
+        _cpu_pack_nvfp4(wrapped)
     else:
-        with pytest.raises(RuntimeError, match="LoRA merge is unsupported after NVFP4 weight quantization"):
-            LoRAPipeline.merge_lora_weights(pipeline)
+        wrapped.merge_lora_weights(accumulate=True)
+        _cpu_pack_nvfp4(wrapped)
+    added = first if entrypoint == "merge" else second
+    _assert_cpu_packed_weight(wrapped, expected + added[1] @ added[0], False)
+    assert len(wrapped._merged_loras) == 2
+    torch.testing.assert_close(wrapped.cpu_weight, original, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("retain", [False, True])
+def test_disjoint_nvfp4_accumulation_then_clear(cpu_nvfp4_lifecycle, retain: bool) -> None:
+    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+
+    pipeline, first_layer, original, first, second = _nvfp4_lifecycle_pipeline(retain)
+    _, second_layer, _, _, _ = _nvfp4_lifecycle_pipeline(retain)
+    model = nn.Module()
+    model.first, model.second = first_layer, second_layer
+    layers = {"first": first_layer, "second": second_layer}
+    pipeline.trainable_transformer_modules = {"transformer": model}
+    pipeline.lora_layers = {"transformer": SimpleNamespace(lora_layers_by_block=lambda: iter([(None, layers)]))}
+    pipeline.lora_adapters = {
+        "a": {"first.lora_A": first[0], "first.lora_B": first[1]},
+        "b": {"second.lora_A": second[0], "second.lora_B": second[1]},
+    }
+    LoRAPipeline.clear_lora_adapter(pipeline)
+    LoRAPipeline.set_lora_adapter(pipeline, "a")
+    LoRAPipeline.set_lora_adapter(pipeline, "b", accumulate=True)
+    _assert_cpu_packed_weight(first_layer, original + first[1] @ first[0], retain)
+    _assert_cpu_packed_weight(second_layer, original + second[1] @ second[0], retain)
+    LoRAPipeline.clear_lora_adapter(pipeline)
+    for layer in layers.values():
+        _assert_cpu_packed_weight(layer, original, retain)
+        assert layer.lora_A is None and layer.lora_B is None
+        assert layer._merged_loras in (None, [])
+
+
+@pytest.mark.parametrize("missing", ["snapshot", "weight_device"])
+def test_nvfp4_invalid_restore_fails_at_layer_boundary_without_mutation(cpu_nvfp4_lifecycle, missing: str) -> None:
+    """Restoration checks its own required state rather than scanning unrelated pipeline modules."""
+    from fastvideo.pipelines.lora_pipeline import LoRAPipeline
+
+    pipeline, wrapped, _, _, _ = _nvfp4_lifecycle_pipeline(False)
+    if missing == "snapshot":
+        wrapped.cpu_weight = None
+    else:
+        wrapped.base_layer.weight_device = None
+    previous_a, previous_b = wrapped.lora_A, wrapped.lora_B
+    with pytest.raises(RuntimeError, match="snapshot|weight_device"):
+        LoRAPipeline.unmerge_lora_weights(pipeline)
+    assert wrapped.base_layer.weight is None and wrapped.merged
+    assert wrapped.lora_A is previous_a and wrapped.lora_B is previous_b
+
+
+def test_nvfp4_restore_uses_placement_marker_without_reading_packed_values(cpu_nvfp4_lifecycle) -> None:
+    _, wrapped, original, _, _ = _nvfp4_lifecycle_pipeline(False)
+    wrapped.base_layer._nvfp4_weight = None
+    wrapped._restore_base_weight()
+    torch.testing.assert_close(wrapped.base_layer.weight, original, rtol=0, atol=0)
+    assert wrapped.base_layer.weight.device == wrapped.base_layer.weight_device.device
+    assert wrapped.base_layer.weight_device.numel() == 0
+    assert "weight_device" not in wrapped.base_layer.state_dict()
+
+
+@pytest.mark.parametrize("requires_grad", [False, True])
+def test_dtensor_lora_restore_preserves_parameter_gradient_flag(tmp_path, requires_grad: bool) -> None:
+    """Replacing a real CPU DTensor parameter preserves freezing or trainability despite the no-grad scope."""
+    import torch.distributed as dist
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.tensor import DTensor, Replicate, distribute_tensor
+    from fastvideo.layers.linear import ReplicatedLinear
+    from fastvideo.layers.lora.linear import get_lora_layer
+
+    if not dist.is_available() or not dist.is_gloo_available():
+        pytest.skip("CPU DTensor restoration requires the Gloo backend")
+    if dist.is_initialized():
+        pytest.skip("This singleton CPU DTensor test must own its process group")
+    dist.init_process_group("gloo", init_method=(tmp_path / "dtensor-store").as_uri(), rank=0, world_size=1)
+    try:
+        mesh = init_device_mesh("cpu", (1, ))
+        original = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+        base = ReplicatedLinear(4, 4, bias=False, params_dtype=torch.float32)
+        base.weight = nn.Parameter(distribute_tensor(original.clone(), mesh, [Replicate()]),
+                                   requires_grad=requires_grad)
+        wrapped = get_lora_layer(base)
+        with torch.no_grad():
+            wrapped.base_layer.weight.add_(5)
+        wrapped.merged = True
+
+        wrapped.unmerge_lora_weights()
+
+        restored = wrapped.base_layer.weight
+        assert isinstance(restored, DTensor)
+        assert restored.requires_grad is requires_grad
+        assert not wrapped.merged
+        torch.testing.assert_close(restored.to_local(), original, rtol=0, atol=0)
+        torch.testing.assert_close(wrapped.cpu_weight.to_local(), original, rtol=0, atol=0)
+        if requires_grad:
+            restored.to_local().sum().backward()
+            assert isinstance(restored.grad, DTensor)
+            torch.testing.assert_close(restored.grad.to_local(), torch.ones_like(original), rtol=0, atol=0)
+        else:
+            assert restored.grad is None
+    finally:
+        dist.destroy_process_group()
+
+
+def test_nvfp4_cpu_snapshot_restore_and_repack_stay_in_offload_scope(cpu_nvfp4_lifecycle, monkeypatch) -> None:
+    from contextlib import contextmanager
+    from fastvideo.pipelines import lora_pipeline
+
+    pipeline, wrapped, _, _, _ = _nvfp4_lifecycle_pipeline(False)
+    events = []
+
+    @contextmanager
+    def scope(module):
+        events.append("enter")
+        try:
+            yield
+        finally:
+            events.append("exit")
+
+    restore = wrapped._restore_base_weight
+
+    def restore_in_scope():
+        assert events[-1] == "enter"
+        events.append("restore")
+        restore()
+
+    def pack_in_scope(module):
+        assert events[-1] == "restore"
+        events.append("pack")
+        _cpu_pack_nvfp4(module)
+
+    monkeypatch.setattr(lora_pipeline, "_get_hook_ctx", scope)
+    monkeypatch.setattr(wrapped, "_restore_base_weight", restore_in_scope)
+    monkeypatch.setattr(nvfp4, "convert_model_to_nvfp4", pack_in_scope)
+    lora_pipeline.LoRAPipeline.unmerge_lora_weights(pipeline)
+    assert events == ["enter", "restore", "pack", "exit"]
+
+
+class _NVFP4LifecycleTransformer(nn.Module):
+
+    def __init__(self, feed_forward: nn.Module) -> None:
+        super().__init__()
+        # A trailing block keeps circular prefetch from targeting its own still-resident FFN.
+        self.transformer_blocks = nn.ModuleList([nn.Sequential(feed_forward), nn.Identity()])
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        for block in self.transformer_blocks:
+            inputs = block(inputs)
+        return inputs
+
+
+@pytest.mark.parametrize("retain", [False, True])
+@pytest.mark.parametrize("layerwise_offload", [False, True])
+@torch.no_grad()
+@torch._dynamo.config.patch(disable=True)
+def test_nvfp4_lora_lifecycle_matches_fresh_quantization_on_blackwell(retain: bool, layerwise_offload: bool) -> None:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
+        pytest.skip("NVFP4 lifecycle numerical validation requires an NVIDIA Blackwell GPU")
+    pytest.importorskip("flashinfer")
+    from fastvideo.hooks.layerwise_offload import enable_layerwise_offload
+    from fastvideo.layers.lora.linear import get_lora_layer
+    from fastvideo.models.dits.minimax_h3 import MiniMaxH3FeedForward
+    from fastvideo.pipelines.lora_pipeline import (LoRAModelLayers, LoRAPipeline,
+                                                   _convert_quantized_weights_after_lora_change, _get_hook_ctx)
+
+    def new_ffn(keep_bf16):
+        previous_dtype = torch.get_default_dtype()
+        torch.set_default_dtype(torch.bfloat16)
+        try:
+            ffn = MiniMaxH3FeedForward(128, 256, quant_config=nvfp4.NVFP4Config(retain_original_weights=keep_bf16),
+                                       prefix="minimax_h3.transformer_blocks.0.ff")
+        finally:
+            torch.set_default_dtype(previous_dtype)
+        return ffn.cuda().eval()
+
+    generator = torch.Generator().manual_seed(20261105)
+
+    def randn(*shape):
+        return torch.randn(shape, generator=generator, dtype=torch.bfloat16).cuda()
+
+    base = {"fc_in": randn(512, 128) * 0.02, "fc_out": randn(128, 256) * 0.02}
+    first = {name: (randn(8, w.shape[1]) * 0.02, randn(w.shape[0], 8)) for name, w in base.items()}
+    inputs = randn(2, 128, 128)
+    second = {name: (randn(4, w.shape[1]) * 0.02, randn(w.shape[0], 4)) for name, w in base.items()}
+    first_weights = {name: w + first[name][1] @ first[name][0] for name, w in base.items()}
+    second_delta = {name: (b @ a) * 0.5 for name, (a, b) in second.items()}
+    accumulated = {name: w + second_delta[name] for name, w in first_weights.items()}
+    replacement = {name: w + second_delta[name] for name, w in base.items()}
+
+    quantized = new_ffn(retain)
+    for name, w in base.items():
+        getattr(quantized, name).weight.copy_(w)
+    transformer = _NVFP4LifecycleTransformer(quantized).eval()
+    block = transformer.transformer_blocks[0]
+    if layerwise_offload:
+        enable_layerwise_offload(transformer)
+    layers = LoRAModelLayers([("transformer_blocks.0", block)])
+    with _get_hook_ctx(block):
+        for name in base:
+            wrapped = get_lora_layer(getattr(quantized, name))
+            setattr(quantized, name, wrapped)
+            wrapped.set_lora_weights(*first[name])
+            layers.add_lora_layer("transformer_blocks.0", name, wrapped)
+        _convert_quantized_weights_after_lora_change([block])
+
+    def assert_state(weights, *, unmerged=False):
+        reference = new_ffn(True)
+        for name, w in weights.items():
+            getattr(reference, name).weight.copy_(w)
+        nvfp4.convert_model_to_nvfp4(reference)
+
+        with _get_hook_ctx(block):
+            for name, w in weights.items():
+                layer = getattr(quantized, name).base_layer
+                expected = getattr(reference, name)
+                if retain:
+                    torch.testing.assert_close(layer.weight, w, rtol=0, atol=0)
+                else:
+                    assert layer.weight is None
+                    assert layer.weight_device.device == inputs.device
+                for key in nvfp4.H3_NVFP4_DIT_BUFFER_NAMES:
+                    actual_buffer, expected_buffer = getattr(layer, key), getattr(expected, key)
+                    assert actual_buffer.dtype == expected_buffer.dtype, f"{name}.{key}"
+                    torch.testing.assert_close(actual_buffer.float(), expected_buffer.float(), rtol=0, atol=0,
+                                               msg=lambda message: f"{name}.{key}: {message}")
+
+        def project(name, x, *, packed):
+            out = getattr(reference, name)(x)[0] if packed else nn.functional.linear(x, weights[name])
+            if unmerged:
+                a, b = first[name]
+                out = out + x @ a.T @ b.T
+            return out
+
+        def expected_output(*, packed):
+            value, gate = project("fc_in", inputs, packed=packed).chunk(2, dim=-1)
+            return project("fc_out", value * nn.functional.silu(gate), packed=packed)
+
+        actual = transformer(inputs)
+        packed_output = expected_output(packed=True)
+        bf16_output = expected_output(packed=False)
+        assert torch.isfinite(actual).all() and torch.isfinite(bf16_output).all()
+        torch.testing.assert_close(actual, packed_output)
+        assert ((actual.float() - bf16_output.float()).norm() <
+                (actual.float() - bf16_base_output.float()).norm())
+
+    base_value, base_gate = nn.functional.linear(inputs, base["fc_in"]).chunk(2, dim=-1)
+    bf16_base_output = nn.functional.linear(base_value * nn.functional.silu(base_gate), base["fc_out"])
+    pipeline = SimpleNamespace(lora_layers={"transformer": layers},
+                               trainable_transformer_modules={"transformer": transformer})
+
+    assert_state(first_weights)
+    LoRAPipeline.unmerge_lora_weights(pipeline)
+    assert_state(base, unmerged=True)
+    LoRAPipeline.merge_lora_weights(pipeline)
+    assert_state(first_weights)
+    for add_to_existing, expected in [(True, accumulated), (False, replacement)]:
+        with _get_hook_ctx(block):
+            for name, factors in second.items():
+                getattr(quantized, name).set_lora_weights(*factors, strength=0.5, accumulate=add_to_existing)
+            _convert_quantized_weights_after_lora_change([block])
+        assert_state(expected)

fastvideo/tests/ops/quantization/test_nvfp4_purge.py

@@ -71,11 +71,16 @@ def test_prefix_classification_sanity() -> None:
 
 def test_auto_purges_always_fp4_and_retains_refine_only(caplog) -> None:
     model = _model(retain=None)
+    original_device = model.always_fp4.weight.device
     with caplog.at_level(logging.INFO):
         nv.convert_model_to_nvfp4(model)
     assert model.always_fp4.weight is None
     assert model.refine_only.weight is not None
     assert model.always_fp4._nvfp4_weight is not None
+    assert model.always_fp4.weight_device.device == original_device
+    assert model.always_fp4.weight_device.untyped_storage().nbytes() == 0
+    assert "weight_device" not in model.always_fp4.state_dict()
+    assert not hasattr(model.refine_only, "weight_device")
     receipt = [r.message for r in caplog.records if "NVFP4 weight purge receipt" in r.message]
     assert receipt and "purged 1" in receipt[0] and "retained 1" in receipt[0]
 
@@ -85,6 +90,36 @@ def test_retain_true_keeps_everything() -> None:
     nv.convert_model_to_nvfp4(model)
     assert model.always_fp4.weight is not None
     assert model.refine_only.weight is not None
+    assert not hasattr(model.always_fp4, "weight_device")
+
+
+def test_purged_weight_placement_follows_module_moves() -> None:
+    model = _model(retain=False)
+    nv.convert_model_to_nvfp4(model)
+    model.to("meta")
+    assert model.always_fp4.weight_device.device.type == "meta"
+    assert model.always_fp4.weight_device.numel() == 0
+    model.to_empty(device="cpu")
+    assert model.always_fp4.weight_device.device.type == "cpu"
+    assert model.always_fp4.weight_device.untyped_storage().nbytes() == 0
+    assert "weight_device" not in model.always_fp4.state_dict()
+
+
+@pytest.mark.parametrize("offload_buffers", [False, True])
+def test_empty_weight_device_participates_in_layerwise_offload(env_overrides, offload_buffers: bool) -> None:
+    import fastvideo.envs as envs
+    from fastvideo.hooks.layerwise_offload import _offload_tensors
+
+    model = _model(retain=False)
+    nv.convert_model_to_nvfp4(model)
+    model.register_buffer("large_payload", torch.empty(1 << 20, dtype=torch.uint8))
+    model.register_buffer("small_cache", torch.ones(4))
+    env_overrides.enter_context(envs.FASTVIDEO_LAYERWISE_OFFLOAD_BUFFERS.override(offload_buffers))
+    selected = dict(_offload_tensors(model))
+    assert selected["always_fp4.weight_device"] is model.always_fp4.weight_device
+    assert ("large_payload" in selected) is offload_buffers
+    assert "small_cache" not in selected
+    assert dict(_offload_tensors(model, selected))["always_fp4.weight_device"] is model.always_fp4.weight_device
 
 
 def test_retained_weight_lora_merge_reaches_nvfp4_buffers(monkeypatch) -> None:

fastvideo/tests/pipelines/test_lora_adapter_switch.py

@@ -113,7 +113,7 @@ def test_clear_restores_base_weights_after_a_change_that_failed_part_way(setup)
     pipeline.lora_adapters["b"] = {**_adapter(["first"], 2), "second.lora_A": torch.randn(RANK, DIM + 1),
                                    "second.lora_B": torch.randn(DIM, RANK)}
 
-    with pytest.raises(RuntimeError):
+    with pytest.raises(ValueError, match="factor dimensions"):
         pipeline.set_lora_adapter("b", strength=1.0)
     pipeline.clear_lora_adapter()
 
@@ -158,7 +158,7 @@ def test_clear_after_a_partial_failure_in_a_block_mixing_side_branch_and_merged_
     pipe._convert_one_transformer = convert_one  # type: ignore[method-assign]
     pipe.set_lora_adapter("a", strength=1.0)
     assert transformer.packed.side_branch and transformer.dense.merged
-    with pytest.raises(RuntimeError):
+    with pytest.raises(ValueError, match="factor dimensions"):
         pipe.set_lora_adapter("b", strength=1.0)
 
     pipe.clear_lora_adapter()