489 lines
20 KiB
Python
489 lines
20 KiB
Python
"""
|
||
SVD Quantizer — Full SVD Compression Logic for Checkpoint Studio
|
||
|
||
Provides the canonical SVD compression logic used across EasyLoRAMerger:
|
||
- SVD_SKIP_PATTERNS — Critical layers to skip SVD entirely
|
||
- SVD_ENERGY_TIERS — Sensitive layer patterns with elevated threshold
|
||
- should_skip_svd() — Skip SVD for critical layer types
|
||
- get_svd_threshold() — Tiered energy threshold (0.999 for sensitive layers)
|
||
- apply_svd_to_tensor() — Per-tensor SVD decomposition + reconstruction
|
||
- apply_svd_preprocess() — Batch SVD preprocessing before quantization
|
||
|
||
All functions are stateless pure functions operating on keys, floats, and tensors.
|
||
This module is the single source of truth; musubi_checkpoint_studio.py
|
||
delegates to it — matching the pattern of fp8_quantizer.py and int8_quantizer.py.
|
||
|
||
The patterns are validated across all six supported architectures:
|
||
- Flux (dot-notation: double_blocks.0.img_attn.qkv.weight)
|
||
- Anima (dot-notation: net.blocks.0.adaln_modulation_cross_attn.1.weight)
|
||
- SDXL (diffusers: model.diffusion_model.input_blocks.0.0.weight)
|
||
- SD1.5 (same SDXL format, no label_emb.)
|
||
- Lumina2, Z-Image (same principles apply)
|
||
"""
|
||
|
||
from typing import Dict, Optional
|
||
|
||
import torch
|
||
|
||
try:
|
||
from ..config import (
|
||
SVD_MIN_DIMENSION,
|
||
SVD_SELECTIVE_MIN_DIM,
|
||
SVD_CLEANUP_PARAM_THRESHOLD,
|
||
)
|
||
from ..utils import DeviceManager, cleanup_memory, memory_guard
|
||
except ImportError:
|
||
from config import (
|
||
SVD_MIN_DIMENSION,
|
||
SVD_SELECTIVE_MIN_DIM,
|
||
SVD_CLEANUP_PARAM_THRESHOLD,
|
||
)
|
||
from utils import DeviceManager, cleanup_memory, memory_guard
|
||
|
||
|
||
|
||
# 🔥 SVD IMPROVEMENT: Patterns that should skip SVD compression entirely
|
||
#
|
||
# SVD compresses weight matrices by truncating small singular values.
|
||
# Signal-entry/exit layers (input projections, output projections,
|
||
# modulation layers, norm weights) are catastrophic to compress because:
|
||
#
|
||
# - Input projections (img_in, time_in, txt_in, vector_in, guidance_in,
|
||
# x_embedder, t_embedder): Signal entry — errors compound through ALL
|
||
# downstream layers.
|
||
#
|
||
# - Output projections (final_layer, out): Signal exit — errors directly
|
||
# visible in pixels.
|
||
#
|
||
# - First block (blocks.0, input_blocks.0, diffusion_model.blocks.0):
|
||
# Everything flows through the first block first. Errors here corrupt
|
||
# every subsequent block.
|
||
#
|
||
# - Modulation layers (single_stream_modulation, double_stream_modulation_*):
|
||
# Controls normalization scale/bias per timestep — tiny errors cause
|
||
# large output shifts. These matrices are typically small (~256×256)
|
||
# with uniformly distributed energy, so 0.95 threshold keeps only
|
||
# 30-50% of vectors — disastrous.
|
||
#
|
||
# - LLM Adapter (llm_adapter.*): T5 cross-attention conditioning interface.
|
||
# First block processes raw T5 hidden states before propagation.
|
||
# Errors corrupt all cross-attention conditioning across every DiT block.
|
||
#
|
||
# - Norm weights (_norm.scale, _norm.bias): Control scale/bias of
|
||
# normalization. Tiny errors → large output shifts because norm
|
||
# statistics are computed per-token.
|
||
#
|
||
# Patterns support two key formats (same as FP8_PRESERVE_BF16_PATTERNS):
|
||
# - Diffusers-style: uses '/' path separators (model/diffusion_model/...)
|
||
# - FLUX-style: uses '.' dot notation (double_blocks.0.img_attn...)
|
||
SVD_SKIP_PATTERNS: tuple = (
|
||
# ── Diffusers-style path separators ──
|
||
'/final_layer/',
|
||
'/time_in/',
|
||
'/txt_in/',
|
||
'/vector_in/',
|
||
'/guidance_in/',
|
||
'/norm.',
|
||
'/input_blocks/0/',
|
||
'/middle_block/',
|
||
'/out/',
|
||
|
||
# ── FLUX dot-notation keys ──
|
||
'final_layer.',
|
||
'time_in.',
|
||
'txt_in.',
|
||
'vector_in.',
|
||
'guidance_in.',
|
||
'img_in.',
|
||
'single_stream_modulation.',
|
||
# FLUX modulation tensors use underscore suffixes (_img, _txt), not dots
|
||
'double_stream_modulation_img.',
|
||
'double_stream_modulation_txt.',
|
||
|
||
# ── SDXL/SD1.5 dot-notation keys ──
|
||
# Time embedding projection — SDXL/SD1.5 equivalent of Flux's time_in.
|
||
# Controls how the denoising timestep signal enters the model.
|
||
# Errors here compound through all denoising steps.
|
||
'time_embed.',
|
||
# Label embedding (SDXL only) — conditioning signal entry point.
|
||
# SDXL uses pooled text embeddings from both CLIP models.
|
||
'label_emb.',
|
||
# First input block — the very first convolution processing the
|
||
# noisy latent. Signal entry point — errors propagate everywhere.
|
||
'input_blocks.0.',
|
||
# UNet middle block (bottleneck) — all features pass through.
|
||
# Small number of tensors but critical for global coherence.
|
||
'middle_block.',
|
||
# Output projection — final convolution before VAE decode.
|
||
# Signal exit point — errors directly affect pixel output.
|
||
'out.',
|
||
|
||
# ── SDXL CLIP Text Encoder — embeddings & first encoder layer ──
|
||
# Token/position embedding tables encode semantic vocabulary.
|
||
# SVD compression corrupts the learned token representations,
|
||
# causing text understanding degradation.
|
||
# Matches: conditioner.embedders.0.transformer.text_model.embeddings.*
|
||
# Note: We do NOT add 'conditioner.embedders.' here — that broad pattern
|
||
# matches BOTH CLIP models (embedders.0 ~123M, embedders.1 ~694M) and is
|
||
# redundant with the two specific patterns below which protect sensitive
|
||
# layers (token embeddings + first encoder layer) in BOTH CLIP models.
|
||
'text_model.embeddings.',
|
||
# First text encoder layer — signal entry for text processing.
|
||
# Similar rationale to net.blocks.0. and input_blocks.0.
|
||
# Matches: ...text_model.encoder.layers.0.*
|
||
'text_model.encoder.layers.0.',
|
||
|
||
# ── Anima architecture patterns ──
|
||
# Input patch embedding — the very first layer processing the noisy latent.
|
||
# Equivalent of Flux's img_in. and SDXL's input_blocks.0.
|
||
# Signal entry — errors here propagate through all 35 DiT blocks.
|
||
# Matches: net.x_embedder.proj.weight
|
||
'net.x_embedder.',
|
||
# Time embedding projection — timestep conditioning signal entry.
|
||
# Equivalent of Flux's time_in. and SDXL's time_embed.
|
||
# Errors here compound through all 35 denoising steps.
|
||
# Matches: net.t_embedder.0.weight, net.t_embedder.2.weight
|
||
'net.t_embedder.',
|
||
# First DiT block — signal entry into the transformer core.
|
||
# All embedded signals (latent, timestep, text conditioning) flow
|
||
# through blocks.0 first. Errors here corrupt every subsequent block.
|
||
# Matches: net.blocks.0.adaln_modulation_cross_attn.1.weight
|
||
'net.blocks.0.',
|
||
# Runtime format variant — ComfyUI wraps UNet under diffusion_model.
|
||
# When a baked/converted checkpoint is re-loaded, keys use the
|
||
# diffusion_model. prefix instead of net. — but the same first-block
|
||
# signal entry rationale applies.
|
||
# Matches: diffusion_model.blocks.0.adaln_modulation_cross_attn.1.weight
|
||
'diffusion_model.blocks.0.',
|
||
# Output projection — final linear layer before VAE decode.
|
||
# Equivalent of Flux's final_layer. and SDXL's out.
|
||
# Signal exit — errors directly affect pixel output.
|
||
# Matches: net.final_layer.linear.weight
|
||
'net.final_layer.',
|
||
|
||
# ── LLM Adapter — Anima's T5 cross-attention conditioning interface ──
|
||
# First llm_adapter block — signal entry for T5 cross-attention conditioning.
|
||
# Similar rationale to net.blocks.0.: the first block processes raw T5
|
||
# hidden states before they propagate to downstream blocks. Errors here
|
||
# corrupt all cross-attention conditioning across every DiT block.
|
||
# Matches: net.llm_adapter.blocks.0.cross_attn.k_proj.weight
|
||
'net.llm_adapter.blocks.0.',
|
||
# LLM Adapter input embedding — projects T5 hidden states into the
|
||
# adapter's feature space. Signal entry point.
|
||
# Matches: net.llm_adapter.embed.weight
|
||
'net.llm_adapter.embed.',
|
||
# LLM Adapter output norm — final normalization before the output
|
||
# projection. Should be preserved alongside other norm parameters.
|
||
# Matches: net.llm_adapter.norm.weight
|
||
'net.llm_adapter.norm.',
|
||
# LLM Adapter output projection — projects adapter features back into
|
||
# the main DiT block's feature space. Signal exit point.
|
||
# Matches: net.llm_adapter.out_proj.weight, net.llm_adapter.out_proj.bias
|
||
'net.llm_adapter.out_proj.',
|
||
|
||
# Suffix-based patterns (work for both formats)
|
||
'_norm.scale',
|
||
'_norm.bias',
|
||
)
|
||
|
||
|
||
# 🔥 SVD ENERGY TIERS: Sensitive-but-large layers get a higher threshold
|
||
#
|
||
# Some layers are too large to skip entirely (they would defeat the purpose
|
||
# of SVD compression) but are still signal-critical enough that the default
|
||
# 0.95 threshold is too aggressive. These get 0.999 instead.
|
||
#
|
||
# The tiered threshold only applies to layers NOT matching SVD_SKIP_PATTERNS.
|
||
# If a layer matches BOTH a skip pattern and a sensitive pattern, the skip
|
||
# takes precedence (it's excluded from SVD entirely).
|
||
SVD_ENERGY_TIERS: dict = {
|
||
# Sensitive layer patterns — signal entry/exit points that are large
|
||
# enough to benefit from SVD but need near-full rank preservation.
|
||
'sensitive_patterns': (
|
||
# ── Diffusers-style path separators ──
|
||
'/final_layer/',
|
||
'/time_in/',
|
||
'/txt_in/',
|
||
'/vector_in/',
|
||
'/guidance_in/',
|
||
# ── FLUX dot-notation keys ──
|
||
'final_layer.',
|
||
'time_in.',
|
||
'txt_in.',
|
||
'vector_in.',
|
||
'guidance_in.',
|
||
'img_in.',
|
||
# ── Anima — signal entry ──
|
||
'net.x_embedder.',
|
||
'net.t_embedder.',
|
||
'net.final_layer.',
|
||
# ── Anima — LLM Adapter first block only (signal entry for T5 cross-attn) ──
|
||
# We do NOT set 0.999 for the entire module — only blocks.0 needs elevated
|
||
# threshold; blocks.1–23 can safely use the default 0.95.
|
||
'net.llm_adapter.blocks.0.',
|
||
# ── SDXL/SD1.5 — signal entry/exit ──
|
||
'time_embed.',
|
||
'label_emb.',
|
||
'out.',
|
||
# ── SDXL output block embedding layers (time embedding projections) ──
|
||
# Similar criticality to time_embed. but on the output side.
|
||
'emb_layers.',
|
||
),
|
||
# Threshold for sensitive layers — 0.999 preserves near-full rank
|
||
# while still allowing some compression for truly redundant modes.
|
||
'sensitive_threshold': 0.999,
|
||
}
|
||
|
||
|
||
# ── Public API ────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def should_skip_svd(key: str) -> bool:
|
||
"""Return True if this key should skip SVD compression entirely.
|
||
|
||
Critical signal-entry/exit layers (img_in, time_in, final_layer,
|
||
modulation, norm weights, first block, LLM adapter entry points)
|
||
are excluded from SVD to prevent quality degradation.
|
||
|
||
Patterns are matched case-insensitively as a substring of the key,
|
||
supporting both Diffusers-style (/) and FLUX/SDXL dot-notation (.).
|
||
|
||
Args:
|
||
key: The checkpoint key to check (e.g. ``double_blocks.0.img_attn.qkv.weight``).
|
||
|
||
Returns:
|
||
True if the key matches a SVD skip pattern.
|
||
"""
|
||
key_lower = key.lower()
|
||
for pattern in SVD_SKIP_PATTERNS:
|
||
if pattern in key_lower:
|
||
return True
|
||
return False
|
||
|
||
|
||
def get_svd_threshold(key: str, default_threshold: float = 0.95) -> float:
|
||
"""Return the energy threshold for SVD compression of a given key.
|
||
|
||
Sensitive layers (signal entry/exit projections) get a higher threshold
|
||
(0.999) to preserve near-full rank. All other layers use the user-provided
|
||
default threshold.
|
||
|
||
NOTE: This function is only called for layers that did NOT match
|
||
``should_skip_svd()``. A layer matching both a skip pattern and a
|
||
sensitive pattern is skipped entirely (skip takes precedence).
|
||
|
||
Args:
|
||
key: The checkpoint key to check.
|
||
default_threshold: The user-specified energy threshold (default 0.95).
|
||
|
||
Returns:
|
||
Energy threshold to use for this key (0.999 for sensitive layers,
|
||
``default_threshold`` otherwise).
|
||
"""
|
||
key_lower = key.lower()
|
||
for pattern in SVD_ENERGY_TIERS['sensitive_patterns']:
|
||
if pattern in key_lower:
|
||
return SVD_ENERGY_TIERS['sensitive_threshold']
|
||
return default_threshold
|
||
|
||
|
||
# ── Per-Tensor SVD Application ────────────────────────────────────────────────
|
||
|
||
|
||
def apply_svd_to_tensor(
|
||
tensor: torch.Tensor,
|
||
key: str,
|
||
svd_mode: str,
|
||
energy_threshold: float,
|
||
target_device: Optional[torch.device] = None,
|
||
) -> tuple:
|
||
"""
|
||
Apply SVD compression to a SINGLE weight tensor and return the
|
||
SVD-reconstructed float32 tensor (or the original if not applicable).
|
||
|
||
This is a pure function — the input tensor is not modified.
|
||
|
||
Args:
|
||
tensor: The weight tensor to compress (must be float dtype).
|
||
key: The checkpoint key (e.g. ``double_blocks.0.img_attn.qkv.weight``).
|
||
svd_mode: ``"none"``, ``"selective"``, or ``"full"``.
|
||
energy_threshold: Energy threshold for rank selection (0.0–1.0).
|
||
target_device: Optional GPU device for accelerated SVD.
|
||
|
||
Returns:
|
||
Tuple of ``(reconstructed_tensor, was_compressed, original_params, compressed_params)``.
|
||
If the tensor does not qualify for SVD, returns ``(original_tensor, False, 0, 0)``.
|
||
"""
|
||
# Must be a 2D weight matrix
|
||
is_weight = key.endswith('.weight') and tensor.ndim == 2
|
||
if not is_weight:
|
||
return tensor, False, 0, 0
|
||
|
||
# Skip critical signal-entry/exit layers entirely.
|
||
if should_skip_svd(key):
|
||
return tensor, False, 0, 0
|
||
|
||
out_features, in_features = tensor.shape
|
||
|
||
# Heuristic for selective mode: only compress large matrices
|
||
if svd_mode == "selective" and (out_features <= SVD_SELECTIVE_MIN_DIM or in_features <= SVD_SELECTIVE_MIN_DIM):
|
||
return tensor, False, 0, 0
|
||
|
||
# Global shape guard: skip tiny matrices regardless of mode
|
||
if out_features < SVD_MIN_DIMENSION and in_features < SVD_MIN_DIMENSION:
|
||
return tensor, False, 0, 0
|
||
|
||
# Skip if rank would always be full rank (no compression benefit)
|
||
min_dim = min(out_features, in_features)
|
||
if min_dim == 1:
|
||
return tensor, False, 0, 0
|
||
|
||
# Convert to float32 for SVD (always operates on float32)
|
||
tensor_fp32 = tensor.float()
|
||
|
||
# ── VRAM-aware per-tensor GPU SVD ──────────────────────────────────
|
||
moved_to_gpu_for_svd = False
|
||
if target_device is not None and target_device.type != "cpu":
|
||
free_vram = DeviceManager.get_free_vram(target_device)
|
||
tensor_bytes = tensor_fp32.numel() * tensor_fp32.element_size()
|
||
needed_bytes = tensor_bytes * 3 # SVD needs ~3× tensor size
|
||
if free_vram is not None and free_vram >= needed_bytes:
|
||
if tensor_fp32.device != target_device:
|
||
tensor_fp32 = tensor_fp32.to(target_device)
|
||
moved_to_gpu_for_svd = True
|
||
short_key = key[-48:]
|
||
print(f" ⚡ GPU SVD [{short_key}] "
|
||
f"{out_features}×{in_features} | "
|
||
f"{tensor_bytes / (1024**2):.1f} MB → "
|
||
f"free VRAM {free_vram / (1024**2):.0f} MB")
|
||
elif free_vram is not None:
|
||
short_key = key[-48:]
|
||
print(f" ⚠️ CPU SVD [{short_key}] "
|
||
f"({out_features}×{in_features}, "
|
||
f"need ≥{needed_bytes / (1024**2):.0f} MB, "
|
||
f"free VRAM {free_vram / (1024**2):.0f} MB)")
|
||
|
||
try:
|
||
U, S, Vh = torch.linalg.svd(tensor_fp32, full_matrices=False)
|
||
except Exception as e:
|
||
print(f"[WARNING] SVD failed for {key}: {e}, skipping")
|
||
return tensor, False, 0, 0
|
||
|
||
# Clamp singular values to prevent NaN from numerically negative tiny SVs
|
||
S = S.clamp(min=0.0)
|
||
|
||
# Determine target rank with tiered threshold
|
||
effective_threshold = get_svd_threshold(key, energy_threshold)
|
||
total_energy = torch.sum(S ** 2)
|
||
cumulative = torch.cumsum(S ** 2, dim=0)
|
||
k = torch.searchsorted(cumulative, effective_threshold * total_energy).item() + 1
|
||
k = min(k, len(S))
|
||
k = max(k, 1) # Ensure at least rank 1
|
||
|
||
# Skip if compressed representation would be larger than original
|
||
original_params = out_features * in_features
|
||
compressed_params = out_features * k + k * in_features
|
||
if compressed_params >= original_params:
|
||
return tensor, False, 0, 0
|
||
|
||
# Truncate and reconstruct
|
||
sqrt_Sk = torch.sqrt(S[:k])
|
||
reconstructed = (U[:, :k] * sqrt_Sk) @ (sqrt_Sk[:, None] * Vh[:k, :])
|
||
|
||
# Move back to CPU if we accelerated on GPU
|
||
if moved_to_gpu_for_svd:
|
||
reconstructed = reconstructed.cpu()
|
||
|
||
if svd_mode == "selective" or svd_mode == "full":
|
||
energy_ratio = torch.sum(S[:k]**2) / total_energy
|
||
print(f"[SVD] compressed {key}: {out_features}x{in_features} -> rank {k} (energy {energy_ratio:.3f})")
|
||
|
||
# Delete intermediate variables to free memory (only for large tensors)
|
||
if out_features * in_features > SVD_CLEANUP_PARAM_THRESHOLD:
|
||
cleanup_memory(U, S, Vh, sqrt_Sk, tensor_fp32)
|
||
memory_guard()
|
||
|
||
return reconstructed, True, original_params, compressed_params
|
||
|
||
|
||
# ── Batch SVD Preprocessing ──────────────────────────────────────────────────
|
||
|
||
|
||
def apply_svd_preprocess(
|
||
tensors: Dict[str, torch.Tensor],
|
||
svd_mode: str,
|
||
energy_threshold: float,
|
||
target_device: Optional[torch.device] = None,
|
||
) -> dict:
|
||
"""
|
||
Apply SVD to original float32 tensors IN PLACE before quantization.
|
||
|
||
Each qualifying weight tensor is decomposed and reconstructed at
|
||
reduced rank. The ``tensors`` dict is mutated: qualifying keys are
|
||
replaced with SVD-reconstructed float32 tensors.
|
||
|
||
Because this runs BEFORE the quantization loop, companion scales
|
||
computed during quantization will AUTOMATICALLY reflect the
|
||
SVD-corrected values — fixing the companion scale invalidation bug.
|
||
|
||
Args:
|
||
tensors: Dict of ``{key: tensor}`` to mutate in-place.
|
||
svd_mode: ``"none"``, ``"selective"``, or ``"full"``.
|
||
energy_threshold: Energy threshold for rank selection.
|
||
target_device: Optional GPU device for accelerated SVD.
|
||
|
||
Returns:
|
||
Dict with svd_info keys for reporting:
|
||
- svd_applied: bool
|
||
- compressed_layers: int
|
||
- skipped_layers: int
|
||
- total_original_params: int
|
||
- total_compressed_params: int
|
||
- size_saved_mb: 0 (SVD does not change storage size)
|
||
- avg_rank_ratio: float
|
||
"""
|
||
if svd_mode == "none":
|
||
return {"svd_applied": False, "compressed_layers": 0,
|
||
"skipped_layers": 0, "size_saved_mb": 0}
|
||
|
||
total_original_params = 0
|
||
total_compressed_params = 0
|
||
compressed_layers = 0
|
||
skipped_layers = 0
|
||
|
||
for key in list(tensors.keys()):
|
||
tensor = tensors[key]
|
||
result, was_compressed, orig_p, comp_p = apply_svd_to_tensor(
|
||
tensor, key, svd_mode, energy_threshold,
|
||
target_device=target_device,
|
||
)
|
||
if was_compressed:
|
||
# Replace with SVD-reconstructed float32 tensor
|
||
tensors[key] = result
|
||
compressed_layers += 1
|
||
total_original_params += orig_p
|
||
total_compressed_params += comp_p
|
||
else:
|
||
skipped_layers += 1
|
||
|
||
# Compute average rank reduction ratio
|
||
avg_rank_ratio = 0.0
|
||
if compressed_layers > 0 and total_original_params > 0:
|
||
# compressed_params is low-rank: out*k + k*in
|
||
# original_params is full: out*in
|
||
# rank ratio ≈ compressed_params / original_params
|
||
avg_rank_ratio = total_compressed_params / total_original_params
|
||
|
||
return {
|
||
"svd_applied": compressed_layers > 0,
|
||
"compressed_layers": compressed_layers,
|
||
"skipped_layers": skipped_layers,
|
||
"total_original_params": total_original_params,
|
||
"total_compressed_params": total_compressed_params,
|
||
# SVD does NOT reduce file size when stored as INT8/FP8 full matrix.
|
||
# Reporting 0 avoids misleading "498 MB saved" messages.
|
||
# The rank_ratio shows the actual compression achieved.
|
||
"size_saved_mb": 0,
|
||
"avg_rank_ratio": avg_rank_ratio,
|
||
}
|