[feat] Add streaming and GPU-accelerated LoRA extraction

.buildkite/scripts/lanes/lora_extraction.sh

@@ -2,4 +2,4 @@
 # Canonical Slurm CI selection for the LoRA-extraction lane.
 set -euo pipefail
 
-exec pytest ./fastvideo/tests/lora_extraction/test_lora_extraction.py -vs
+exec pytest ./fastvideo/tests/lora_extraction/ -vs

docs/training/finetune.md

@@ -110,7 +110,7 @@ Key differences from full finetune:
 
 ## LoRA Extraction and Merging
 
-FastVideo provides tools to extract LoRA adapters from finetuned models and merge them back.
+FastVideo provides generic runtime adapter extraction and retains a legacy merger for its previously supported adapter layouts.
 
 ### Extract LoRA Adapter
 
@@ -121,25 +121,41 @@ python scripts/lora_extraction/extract_lora.py \
   --base Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
   --ft path/to/your/finetuned_model \
   --out adapter_r32.safetensors \
-  --rank 32
+  --rank 32 \
+  --exact-tensor-pattern '^condition_embedder\.' \
+  --exact-tensor-pattern '^proj_out\.weight$'
 ```
 
 | Argument | Description |
 |----------|-------------|
-| `--base` | Base model (HuggingFace ID or local path) |
-| `--ft` | Finetuned model path |
-| `--out` | Output adapter file (.safetensors) |
-| `--rank` | LoRA rank (16, 32, 64, 128) |
+| `--base` | Base model (Hugging Face ID or local path) |
+| `--ft` | Fine-tuned model (Hugging Face ID or local path) |
+| `--out` | Output adapter file (`.safetensors`) |
+| `--rank` | Requested LoRA rank (for example, 16, 32, 64, or 128) |
 | `--full-rank` | Extract full-rank adapter (optional) |
+| `--min-delta` | Omit tensors whose maximum absolute FP32 delta is at or below the threshold (default: `1e-8`) |
+| `--load-mode` | `auto` (indexed, then pipeline fallback), `indexed`, or `pipeline` |
+| `--device` | SVD device, such as `cpu` or `cuda:0` |
+| `--svd-method` | Exact or randomized SVD |
+| `--factor-dtype` | Storage dtype for the low-rank factors |
+| `--dense-dtype` | Storage dtype for exact `.diff`/`.diff_b`/`.diff_param` payloads (default: `float32`) |
+| `--exact-tensor-pattern` | Repeatable regex for matrices the target runtime cannot load as LoRA factors |
 
-### Merge LoRA Adapter
+For large checkpoints, indexed loading streams one transformer tensor pair at a time and downloads only `transformer/*`. The extractor also preserves changed norms, biases, and standalone parameters as exact deltas, and fine-tuned-only parameters as `.set_weight` or `.set_param`. Matrix selection is runtime-agnostic, so full-finetune extraction must keep runtime-unsupported matrices exact, as the Wan example does. See the [LoRA utilities](../utilities/lora.md) for the GPU/randomized-SVD command, resume options, and accuracy controls.
 
-Merge an adapter back into a base model:
+Mixed low-rank/dense adapters from the generic extractor must be supplied when constructing FastVideo through
+`ComponentConfig(lora_path=...)`; their dense payload cannot be swapped later with `set_lora_adapter`. The legacy
+offline merger below is not part of this extraction workflow.
+
+### Legacy Merge LoRA Adapter
+
+The command below documents the pre-existing merger for adapters it already supports. Do not pass a mixed adapter from
+the generic extractor to it: the legacy merger does not apply exact dense or replacement payloads.
 
 ```bash
 python scripts/lora_extraction/merge_lora.py \
   --base Wan-AI/Wan2.1-T2V-1.3B-Diffusers \
-  --adapter adapter_r32.safetensors \
+  --adapter legacy_factor_only_adapter.safetensors \
   --ft path/to/your/finetuned_model \
   --output merged_model
 ```

docs/utilities/lora.md

@@ -1,6 +1,6 @@
 # LoRA Extraction and Merging
 
-Tools for extracting and merging LoRA adapters for FastVideo models.
+Generic runtime adapter extraction plus the existing legacy merge utilities for FastVideo models.
 
 ## Extract LoRA Adapter
 
@@ -9,30 +9,75 @@ python scripts/lora_extraction/extract_lora.py \
   --base Wan-AI/Wan2.2-TI2V-5B-Diffusers \
   --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
   --out adapter_r32.safetensors \
-  --rank 32
+  --rank 32 \
+  --exact-tensor-pattern '^condition_embedder\.' \
+  --exact-tensor-pattern '^proj_out\.weight$'
 ```
 
-**Options:**
+The extractor is runtime-agnostic and cannot determine from checkpoint tensors whether the target runtime
+wraps a given matrix as a LoRA layer. Use `--exact-tensor-pattern` for changed matrices that the runtime does not wrap;
+the extractor preserves them as exact `.diff` tensors instead of emitting factors that the runtime cannot apply. The
+Wan patterns above cover its excluded condition embedders and its unwrapped output projection.
+
+Exact CPU SVD remains the default. For a large transformer, stream its indexed safetensors and factorize on a GPU:
+
+```bash
+python scripts/lora_extraction/extract_lora.py \
+  --base <base-model-or-path> \
+  --ft <finetuned-model-or-path> \
+  --out adapter_r64.safetensors \
+  --rank 64 \
+  --load-mode indexed \
+  --device cuda:0 \
+  --svd-method randomized \
+  --randomized-q 320 \
+  --niter 4 \
+  --factor-dtype float16 \
+  --dense-dtype float32 \
+  --replacement-dtype source
+```
+
+`--load-mode indexed` downloads only `transformer/*` for a Hugging Face model and reads one base/fine-tuned tensor pair at a time. The default `auto` mode tries indexed loading first and falls back to the legacy FastVideo pipeline loader; `pipeline` selects the legacy loader directly.
+
+Important options:
+
+- `--base`, `--ft`: Hugging Face model IDs or local paths.
+- `--rank`, `--full-rank`: truncated or full factorization rank.
+- `--min-delta`: omit tensors whose maximum absolute FP32 delta is at or below this threshold (default: `1e-8`).
+- `--device`: factorization device, such as `cpu` or `cuda:0`.
+- `--svd-method`: `exact` or `randomized`.
+- `--randomized-q`, `--niter`, `--seed`: randomized SVD accuracy and reproducibility.
+- `--factor-dtype`, `--dense-dtype`, `--replacement-dtype`: adapter storage precision. Exact dense deltas default to `float32`.
+- `--exact-tensor-pattern`: repeatable regex for a matrix that should remain an exact dense delta.
+- `--base-revision`, `--ft-revision`: pin Hugging Face inputs in indexed mode; revisions are rejected for local paths and pipeline loading.
+- `--work-dir`, `--resume`: resume an interrupted streaming extraction. Scratch is written to an
+  output-specific namespace under `fastvideo-lora-extract/`, and only that namespace is cleaned up. Resume requires indexed
+  safetensors and validates both checkpoints' index/shard fingerprints before reusing partial results.
+
+The adapter retains changes that do not fit a low-rank product: `.diff` and `.diff_b` hold exact additive weight/bias deltas, `.diff_param` handles standalone parameters such as `scale_shift_table`, and `.set_weight`/`.set_param` hold parameters absent from the base checkpoint. Bit-identical parameters are omitted. The extractor writes an adjacent `*.report.json` with tensor counts, settings, and reconstruction residuals.
+
+For the validated MiniMax-H3 rank-64 command, including its exact-boundary patterns, see [`scripts/lora_extraction/README.md`](https://github.com/hao-ai-lab/FastVideo/blob/main/scripts/lora_extraction/README.md).
+
+Mixed low-rank/dense adapters produced by the generic extractor must be supplied when constructing FastVideo through
+`ComponentConfig(lora_path=...)`; their dense payload cannot be swapped later with `set_lora_adapter`. The legacy
+offline merger below retains its existing scope and is not part of this extraction workflow.
 
-- `--base`: Base model (HuggingFace ID or local path)
-- `--ft`: Fine-tuned model (HuggingFace ID or local path)
-- `--out`: Output adapter file
-- `--rank`: LoRA rank (16, 32, 64, 128)
-- `--full-rank`: Extract full-rank adapter (optional)
+## Legacy Merge Adapter
 
-## Merge Adapter
+The command below documents the pre-existing merger for adapters it already supports. Do not pass a mixed adapter from
+the generic extractor to it: the legacy merger does not apply the adapter's exact dense or replacement payloads.
 
 ```bash
 python scripts/lora_extraction/merge_lora.py \
   --base Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --adapter adapter_r32.safetensors \
+  --adapter legacy_factor_only_adapter.safetensors \
   --ft FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers \
   --output merged_model
 ```
 
 **Options:**
 
-- `--base`: Base model (HuggingFace ID or local path)
+- `--base`: Base model (Hugging Face ID or local path)
 - `--adapter`: LoRA adapter file (.safetensors)
 - `--ft`: Fine-tuned model (for configuration)
 - `--output`: Output directory

fastvideo/models/loader/fsdp_load.py

@@ -241,11 +241,17 @@ def maybe_load_fsdp_model(
 
     weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
     param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
-    dense_lora_patch = DenseLoRAPatch.from_adapter(
-        lora_path,
-        param_names_mapping_fn,
-        strength=lora_strength,
-    )
+    dense_lora_patch = None
+    if lora_path:
+        # LoRA-specific mappings are optional and must not become a dependency of ordinary model loading.
+        lora_mapping = getattr(model, "lora_param_names_mapping", None)
+        lora_param_names_mapping_fn = get_param_names_mapping(lora_mapping) if lora_mapping else None
+        dense_lora_patch = DenseLoRAPatch.from_adapter(
+            lora_path,
+            param_names_mapping_fn,
+            lora_param_names_mapping=lora_param_names_mapping_fn,
+            strength=lora_strength,
+        )
     if dense_lora_patch is not None:
         # H3's compression gate is created only by the VSA attention backend. Loading a
         # VSA student under dense attention would otherwise warn about 50 unmatched
@@ -695,7 +701,7 @@ def load_model_from_full_model_state_dict(
             target_dtype = dtype_selector(new_param_name, param_dtype)
         if adapter_value is not None:
             if tuple(adapter_value.shape) != tuple(meta_sharded_param.shape):
-                raise ValueError(f"LoRA set_weight for {new_param_name} has shape {tuple(adapter_value.shape)}, "
+                raise ValueError(f"LoRA replacement for {new_param_name} has shape {tuple(adapter_value.shape)}, "
                                  f"but the parameter is {tuple(meta_sharded_param.shape)}")
             full_tensor = adapter_value.to(device=device, dtype=target_dtype)
             if not hasattr(meta_sharded_param, "device_mesh"):

fastvideo/models/loader/lora_patch.py

@@ -6,16 +6,17 @@
 rank. Distilled video checkpoints break both often enough that dropping whatever does
 not fit silently loses real signal.
 
-Two payload kinds cover the gap, named after the convention ComfyUI's loader already
-reads so one file works in both places:
+Two payload kinds cover the gap. The weight/bias spellings follow conventions used by
+ComfyUI; the ``*_param`` spellings extend them to standalone FastVideo parameters:
 
-``<module>.diff`` / ``<module>.diff_b``
+``<module>.diff`` / ``<module>.diff_b`` / ``<parameter>.diff_param``
     An exact additive delta for a parameter the base model has. Used where a rank-``r``
     factorization buys nothing or cannot be formed at all -- RMSNorm vectors, biases,
-    and matrices whose smaller dimension is already at or below the rank that would be
-    chosen. Factoring a length-``n`` vector into rank ``r`` costs ``r(1 + n) > n``.
+    standalone parameters such as ``scale_shift_table``, and matrices whose smaller
+    dimension is already at or below the rank that would be chosen. Factoring a
+    length-``n`` vector into rank ``r`` costs ``r(1 + n) > n``.
 
-``<module>.set_weight``
+``<module>.set_weight`` / ``<parameter>.set_param``
     A whole parameter the base model does not carry, so no delta is expressible. MiniMax
     H3's VSA ``to_gate_compress`` is the case that motivated this: it exists only under
     the sparse-attention backend, and :func:`load_model_from_full_model_state_dict`
@@ -49,17 +50,11 @@
 
 logger = init_logger(__name__)
 
-# Suffix -> the parameter suffix it targets. ``.diff``/``.diff_b`` are additive,
-# ``.set_weight`` replaces. Ordered longest-first so ``.diff_b`` is tested before
-# ``.diff`` would match a truncated key.
-ADDITIVE_SUFFIXES: dict[str, str] = {".diff_b": ".bias", ".diff": ".weight"}
-REPLACEMENT_SUFFIXES: dict[str, str] = {".set_weight": ".weight"}
-
-# Recognized elsewhere in an adapter and deliberately not our business: the low-rank
-# half, which ``LoRAPipeline`` merges through the wrapped-module path.
-_LOW_RANK_MARKERS = (".lora_A", ".lora_B", ".lora_up", ".lora_down", ".lora_alpha", ".lora_rank", ".alpha",
-                     ".dora_scale")
-
+# Suffix -> the parameter suffix it targets. An empty target suffix preserves the
+# full parameter name for standalone nn.Parameters. Ordered longest-first so the
+# generic spellings are tested before the shorter weight/bias spellings.
+ADDITIVE_SUFFIXES: dict[str, str] = {".diff_param": "", ".diff_b": ".bias", ".diff": ".weight"}
+REPLACEMENT_SUFFIXES: dict[str, str] = {".set_weight": ".weight", ".set_param": ""}
 
 # One low-rank pair has many spellings. PEFT writes ``.lora_A.weight``, and interposes
 # the adapter's name when it is not the default (``.lora_A.default.weight``); kohya and
@@ -126,14 +121,14 @@ def from_adapter(
         lora_path: str | None,
         param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
         *,
+        lora_param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
         strength: float = 1.0,
     ) -> DenseLoRAPatch | None:
         """Build a patch from an adapter, or ``None`` when it carries no dense payload.
 
-        ``param_names_mapping`` is the same callable the checkpoint loader uses, so
-        adapter keys are resolved into the model's own parameter names by the identical
-        rules -- an adapter written against the published checkpoint layout needs no
-        separate conversion table.
+        ``lora_param_names_mapping`` first translates adapter-specific official names
+        into the published checkpoint layout. ``param_names_mapping`` then resolves that
+        layout into the model's parameter names, matching the low-rank loader's order.
         """
         if not lora_path:
             return None
@@ -150,7 +145,7 @@ def from_adapter(
         for path in files:
             with safe_open(path, framework="pt") as handle:
                 for key in handle.keys():
-                    resolved = _resolve(key, param_names_mapping)
+                    resolved = _resolve(key, lora_param_names_mapping, param_names_mapping)
                     if resolved is None:
                         continue
                     target, kind = resolved
@@ -163,7 +158,7 @@ def from_adapter(
         if not additive and not replacement:
             return None
         logger.info(
-            "LoRA adapter %s carries a dense payload: %d additive (.diff/.diff_b), %d replacement (.set_weight)",
+            "LoRA adapter %s carries a dense payload: %d additive, %d replacement parameters",
             lora_path, len(additive), len(replacement))
         return cls(files, additive, replacement, strength)
 
@@ -234,31 +229,49 @@ def _read(self, entry: tuple[str, str]) -> torch.Tensor:
 
 def _resolve(
     key: str,
+    lora_param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None,
     param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None,
 ) -> tuple[str, str] | None:
     """Map an adapter key to ``(model parameter name, "add" | "set")``.
 
     Returns ``None`` for anything that is not a dense payload key, which includes every
     low-rank factor -- those belong to ``LoRAPipeline``, not here.
     """
-    if any(marker in key for marker in _LOW_RANK_MARKERS):
-        return None
+    # A terminal dense suffix is authoritative even when an ordinary module name
+    # contains text such as ``.alpha`` or ``.lora_A``.
     for suffix, param_suffix in ADDITIVE_SUFFIXES.items():
         if key.endswith(suffix):
-            return _map_name(key[:-len(suffix)] + param_suffix, param_names_mapping, key), "add"
+            return _map_name(
+                key[:-len(suffix)] + param_suffix,
+                lora_param_names_mapping,
+                param_names_mapping,
+                key,
+            ), "add"
     for suffix, param_suffix in REPLACEMENT_SUFFIXES.items():
         if key.endswith(suffix):
-            return _map_name(key[:-len(suffix)] + param_suffix, param_names_mapping, key), "set"
+            return _map_name(
+                key[:-len(suffix)] + param_suffix,
+                lora_param_names_mapping,
+                param_names_mapping,
+                key,
+            ), "set"
     return None
 
 
 def _map_name(
     param_name: str,
+    lora_param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None,
     param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None,
     source_key: str,
 ) -> str:
-    """Run the checkpoint loader's own renaming rules over a resolved parameter name."""
+    """Run the low-rank loader's two-stage renaming rules over a dense parameter."""
     param_name = param_name.replace("diffusion_model.", "")
+    if lora_param_names_mapping is not None:
+        param_name, merge_index, _ = lora_param_names_mapping(param_name)
+        if merge_index is not None:
+            raise NotImplementedError(f"LoRA dense key {source_key} resolves to a fused parameter during the "
+                                      "adapter-specific mapping; whole-tensor payloads for fused parameters "
+                                      "are not supported")
     if param_names_mapping is None:
         return param_name
     mapped, merge_index, _ = param_names_mapping(param_name)

fastvideo/pipelines/lora_pipeline.py

@@ -346,12 +346,12 @@ def set_lora_adapter(self,
         if not self._setting_constructor_adapter:
             if self._constructor_dense_lora_path is not None:
                 raise RuntimeError(
-                    "The active LoRA contains constructor-time .diff/.set_weight payload. "
+                    "The active LoRA contains a constructor-time dense additive/replacement payload. "
                     "Changing its adapter or strength at runtime would leave that dense payload stale; "
                     "create a new VideoGenerator with ComponentConfig(lora_path=..., lora_strength=...).")
             if requested_path is not None and DenseLoRAPatch.from_adapter(requested_path) is not None:
                 raise RuntimeError(
-                    "Adapters containing .diff/.set_weight payload must be supplied when VideoGenerator is "
+                    "Adapters containing dense additive/replacement payloads must be supplied when VideoGenerator is "
                     "constructed with ComponentConfig(lora_path=..., lora_strength=...).")
 
         if lora_nickname not in self.lora_adapters and lora_path is None: