[feat] Add streaming and GPU-accelerated LoRA extraction

[shaoxiongduan]

Purpose

Add a streaming, GPU-accelerated path to scripts/lora_extraction/extract_lora.py for large video transformers while preserving the existing exact CPU workflow and mixed LoRA/dense adapter format.

Previously, extraction materialized both transformer state dictionaries in host memory and performed full SVD on CPU. This is prohibitively expensive for models such as MiniMax-H3. The updated extractor streams indexed tensors, supports exact or randomized factorization on a selected device, and adds controls for accuracy, storage precision, resumability, and adapter validation.

Changes

• Stream indexed safetensors one base/fine-tuned tensor pair at a time.
• Download only transformer/* when resolving Hugging Face repositories.
• Preserve the existing loading behavior through --load-mode auto:
  • try indexed loading first;
  • fall back to the legacy FastVideo pipeline loader;
  • allow either path to be selected explicitly.
• Add configurable CPU/GPU factorization:
  • --device
  • --svd-method exact|randomized
  • --randomized-q
  • --oversample
  • --niter
  • --seed
• Add storage dtype controls for:
  • low-rank factors;
  • exact .diff / .diff_b / .diff_param payloads;
  • fine-tuned-only .set_weight / .set_param parameters.
• Add repeatable --exact-tensor-pattern rules for matrices that should remain exact dense deltas instead of being rank-truncated.
• Document that matrix selection is runtime-agnostic and keep Wan condition embedders and its unwrapped output projection exact in the validated extraction command.
• Add resumable per-tensor extraction through --work-dir and --resume.
• Validate checkpoint key sets, tensor shapes, generated adapter keys, and output factor shapes.
• Preserve changed standalone parameters such as Wan2.2 scale_shift_table through .diff_param/.set_param payloads instead of silently dropping them.
• Write an adjacent extraction report containing settings, tensor counts, and reconstruction residuals.
• Omit unnecessary per-layer lora_rank and lora_alpha tensors when alpha == rank, which is already the loader default.
• Retain the existing mixed-adapter behavior:
  • changed matrices → lora_B @ lora_A;
  • changed base weights/biases → .diff / .diff_b;
  • changed standalone parameters such as scale_shift_table → .diff_param;
  • fine-tuned-only weights such as VSA gates → .set_weight;
  • other fine-tuned-only parameters → .set_param;
  • bit-identical parameters → omitted.
• Update the LoRA extraction documentation and add CPU/GPU unit coverage.

Exact CPU SVD remains the default for backward compatibility. GPU and randomized SVD are opt-in.

Test Plan

source /mnt/lustre/vlm-s4duan/FastVideo/.venv/bin/activate

# CPU unit coverage
PYTHONPATH="$PWD" pytest \
  --confcutdir=fastvideo/tests/lora_extraction \
  fastvideo/tests/lora_extraction/test_streaming_lora_extraction.py -q

# Run the same test on a GPU node to activate CUDA coverage
PYTHONPATH="$PWD" pytest \
  --confcutdir=fastvideo/tests/lora_extraction \
  fastvideo/tests/lora_extraction/test_streaming_lora_extraction.py -q

# Existing dense LoRA loader coverage
PYTHONPATH="$PWD" pytest \
  fastvideo/tests/loader/test_lora_patch.py -q

# Wan2.2 extraction through the real FastVideo DMD LoRA loader on GPU
PYTHONPATH="$PWD" pytest \
  --confcutdir=fastvideo/tests/lora_extraction \
  fastvideo/tests/lora_extraction/test_lora_extraction.py \
  -q -s

# Full repository pre-commit suite
pre-commit run --all-files

A full MiniMax-H3 extraction was also run on one GB200:

python scripts/lora_extraction/extract_lora.py \
  --base /mnt/lustre/vlm-s4duan/models/MiniMax-H3 \
  --ft /mnt/lustre/vlm-s4duan/models/FastVideo-FastH3-8-step-Preview-v1-VSA-DataFree-full-local \
  --out /tmp/fasth3-generic-r64/adapter_model.safetensors \
  --rank 64 \
  --min-delta 0 \
  --load-mode indexed \
  --device cuda:0 \
  --svd-method randomized \
  --randomized-q 320 \
  --niter 4 \
  --seed 42 \
  --factor-dtype float16 \
  --dense-dtype float32 \
  --replacement-dtype source \
  --exact-tensor-pattern '^audio_proj_(in|out)\.weight$' \
  --exact-tensor-pattern '^context_embedder\.weight$' \
  --exact-tensor-pattern '^proj_(in|out)\.weight$' \
  --exact-tensor-pattern '^time_embedder\.'

The extracted adapter was then loaded through the MiniMax-H3 LoRA inference path and used to generate an eight-step FastH3 video.

Test Results

Test output

CPU extraction tests:
5 passed, 2 skipped

GPU extraction tests:
7 passed

Dense LoRA loader tests:
30 passed

Wan2.2 GPU extraction and real-loader integration test:
1 passed in 250.23s

Wan2.2 adapter application:
525 exact dense parameters fully applied
300 LoRA layers converted and applied
0 unmatched factor keys
WanDMDPipeline selected

Pre-commit (`pre-commit run --all-files`):
yapf passed
ruff passed
codespell passed
PyMarkdown passed
actionlint passed
mypy passed
filename check passed
suggestion check passed

Full MiniMax-H3 extraction:

LoRA layers:             362
Exact additive tensors:   78
VSA replacement tensors:  50
Unchanged tensors:       198
Weighted residual:       0.8755630793425291

Comparison against the adapter previously produced by the specialized MiniMax-H3 extractor:

{
  "old_keys": 852,
  "new_keys": 852,
  "key_sets_equal": true,
  "unequal_tensors": 0,
  "max_abs_difference": 0.0
}

End-to-end FastH3 inference:

LoRA dense payload fully applied: 128 parameters
LoRA adapter applied to 362 layers
Video generation completed successfully

The existing Wan2.2 integration example now extracts on cuda:0, retains runtime-unsupported boundary matrices as exact deltas, constructs the real FastVideo WanDMDPipeline, and asserts that all 300 factorized matrices reach a layer.

Checklist

• I ran pre-commit run --all-files and fixed all issues
• I added or updated tests for my changes
• I updated documentation if needed
• I considered GPU memory impact of my changes

For model/pipeline changes, also check:

Not applicable; this PR changes the extraction utility and does not modify model or pipeline numerics.

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=fastcheck-passed
• check-success=full-suite-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[SolitaryThinker]

Requesting changes for several silent correctness failures in the documented extraction and merge paths. I reproduced the namespace mismatch, stale-resume behavior, ignored revisions, and low-precision dense round-trip locally. The focused CPU tests and pre-commit pass, but the current tests do not exercise these production cases; Buildkite fastcheck is also still red across all microscope lanes.

[mergify]

This PR has merge conflicts with the base branch. Please rebase:

git fetch origin main
git rebase origin/main
# Resolve any conflicts, then:
git push --force-with-lease