diff --git a/docs/inference/h3_community_lora.md b/docs/inference/h3_community_lora.md index 4e66903eb..e559b9efa 100644 --- a/docs/inference/h3_community_lora.md +++ b/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 diff --git a/fastvideo/hooks/layerwise_offload.py b/fastvideo/hooks/layerwise_offload.py index 211ea0f3e..e2979002d 100644 --- a/fastvideo/hooks/layerwise_offload.py +++ b/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: diff --git a/fastvideo/layers/lora/linear.py b/fastvideo/layers/lora/linear.py index a65ef28c0..ab793c301 100644 --- a/fastvideo/layers/lora/linear.py +++ b/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): """ diff --git a/fastvideo/layers/quantization/nvfp4_config.py b/fastvideo/layers/quantization/nvfp4_config.py index 033771145..2d082f625 100644 --- a/fastvideo/layers/quantization/nvfp4_config.py +++ b/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: diff --git a/fastvideo/pipelines/lora_pipeline.py b/fastvideo/pipelines/lora_pipeline.py index e482ae651..0c80a2054 100644 --- a/fastvideo/pipelines/lora_pipeline.py +++ b/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, diff --git a/fastvideo/tests/hooks/test_layerwise_offload.py b/fastvideo/tests/hooks/test_layerwise_offload.py index f1505029e..fe72c7a60 100644 --- a/fastvideo/tests/hooks/test_layerwise_offload.py +++ b/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.""" diff --git a/fastvideo/tests/layers/test_lora_side_branch.py b/fastvideo/tests/layers/test_lora_side_branch.py index d1d164492..f64959ad7 100644 --- a/fastvideo/tests/layers/test_lora_side_branch.py +++ b/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 diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py b/fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py index 32b9f84d1..64f78fa66 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_h3_dit_export.py +++ b/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 diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_lora_accumulation.py b/fastvideo/tests/ops/quantization/test_nvfp4_lora_accumulation.py new file mode 100644 index 000000000..f58957d77 --- /dev/null +++ b/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 diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py b/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py index 1f34356ac..11fb2d2eb 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_minimax_h3_wiring.py +++ b/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) diff --git a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py b/fastvideo/tests/ops/quantization/test_nvfp4_purge.py index f44b39309..ad053a07f 100644 --- a/fastvideo/tests/ops/quantization/test_nvfp4_purge.py +++ b/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: diff --git a/fastvideo/tests/pipelines/test_lora_adapter_switch.py b/fastvideo/tests/pipelines/test_lora_adapter_switch.py index 9bac068f9..187e7751f 100644 --- a/fastvideo/tests/pipelines/test_lora_adapter_switch.py +++ b/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()