617 lines
26 KiB
Python
617 lines
26 KiB
Python
"""
|
|
SNR-Quant - ComfyUI nodes for adaptive FP8/ BF16 quantization.
|
|
"""
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
from safetensors.torch import load_file, save_file
|
|
from typing import List, Dict, Tuple, Optional
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
|
|
import folder_paths
|
|
from comfy.utils import ProgressBar
|
|
|
|
|
|
class ModelFragility(Enum):
|
|
"""Model fragility levels."""
|
|
FRAGILE = "fragile"
|
|
MODERATE = "moderate"
|
|
ROBUST = "robust"
|
|
|
|
|
|
@dataclass
|
|
class LayerInfo:
|
|
"""Information about a single layer."""
|
|
name: str
|
|
numel: int
|
|
snr: float
|
|
outlier: float
|
|
|
|
@property
|
|
def bf16_size(self) -> int:
|
|
return self.numel * 2
|
|
|
|
@property
|
|
def fp8_size(self) -> int:
|
|
return self.numel
|
|
|
|
|
|
def compute_snr(weight: torch.Tensor) -> float:
|
|
"""Compute Signal-to-Noise Ratio for FP8 E4M3 quantization."""
|
|
f32_w = weight.float()
|
|
|
|
if torch.isnan(f32_w).any() or torch.isinf(f32_w).any():
|
|
return float('inf')
|
|
|
|
max_abs = torch.max(torch.abs(f32_w)).item()
|
|
if max_abs == 0:
|
|
return float('inf')
|
|
|
|
scale = max_abs / 448.0
|
|
q = (f32_w / scale).clamp(-448, 448)
|
|
q = q.to(torch.float8_e4m3fn).to(f32_w.dtype)
|
|
dq = q * scale
|
|
|
|
signal_pwr = torch.mean(f32_w ** 2).item()
|
|
noise_pwr = torch.mean((f32_w - dq) ** 2).item()
|
|
|
|
if noise_pwr == 0:
|
|
return float('inf')
|
|
|
|
return 10 * np.log10(signal_pwr / noise_pwr)
|
|
|
|
|
|
def compute_outlier_index(weight: torch.Tensor) -> float:
|
|
"""Compute outlier index = max_abs / std."""
|
|
f32_w = weight.float()
|
|
max_abs = torch.max(torch.abs(f32_w)).item()
|
|
std = torch.std(f32_w).item()
|
|
if std == 0:
|
|
return 0.0
|
|
return max_abs / std
|
|
|
|
|
|
def analyze_model(sd: Dict[str, torch.Tensor]) -> Tuple[List[LayerInfo], int, int]:
|
|
"""Analyze all 2D weight layers.
|
|
|
|
Returns:
|
|
layers: List of LayerInfo for quantizable (non-FP8) layers
|
|
total_params: Total params in 2D layers
|
|
existing_fp8_count: Number of already quantized FP8 layers
|
|
"""
|
|
layers = []
|
|
total_params = 0
|
|
existing_fp8_count = 0
|
|
|
|
for name, tensor in sd.items():
|
|
if tensor.dim() == 2:
|
|
numel = tensor.numel()
|
|
total_params += numel
|
|
|
|
# Skip already quantized FP8 layers
|
|
if tensor.dtype == torch.float8_e4m3fn:
|
|
existing_fp8_count += 1
|
|
continue
|
|
|
|
snr = compute_snr(tensor)
|
|
outlier = compute_outlier_index(tensor)
|
|
layers.append(LayerInfo(name=name, numel=numel, snr=snr, outlier=outlier))
|
|
|
|
return layers, total_params, existing_fp8_count
|
|
|
|
|
|
def detect_fragility(layers: List[LayerInfo]) -> Tuple[ModelFragility, float, float]:
|
|
"""Detect model fragility based on SNR distribution."""
|
|
snrs = [l.snr for l in layers]
|
|
snr_range = max(snrs) - min(snrs)
|
|
snr_median = np.median(snrs)
|
|
|
|
if snr_range < 2 and snr_median < 32:
|
|
return ModelFragility.FRAGILE, snr_range, snr_median
|
|
elif snr_range > 5 and snr_median > 34:
|
|
return ModelFragility.ROBUST, snr_range, snr_median
|
|
else:
|
|
return ModelFragility.MODERATE, snr_range, snr_median
|
|
|
|
|
|
class SNRQuantNode:
|
|
"""Quantize model to target size using FP8/ BF16 mixed precision."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("diffusion_models"), {
|
|
"tooltip": "Model from ComfyUI/models/diffusion_models"
|
|
}),
|
|
"target_size_GB": ("FLOAT", {
|
|
"default": 22.0,
|
|
"min": 1.0,
|
|
"step": 0.1,
|
|
"tooltip": "Target model size in GB (e.g., 22.0)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "quantize"
|
|
CATEGORY = "SNR-Quant"
|
|
|
|
def quantize(self, model_name: str, target_size_GB: float):
|
|
# Get paths
|
|
model_path = folder_paths.get_full_path("diffusion_models", model_name)
|
|
output_dir = folder_paths.get_output_directory()
|
|
# Create diffusion_models folder in output if it doesn't exist
|
|
diffusion_models_dir = os.path.join(output_dir, "diffusion_models")
|
|
if not os.path.exists(diffusion_models_dir):
|
|
os.makedirs(diffusion_models_dir)
|
|
# Replace path separators with underscores for flat output filename
|
|
base_name = os.path.splitext(model_name)[0].replace('\\', '_').replace('/', '_')
|
|
# Include target size in filename to distinguish multiple quantizations
|
|
output_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB.safetensors")
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"[SNR-Quant] Auto Quant - FP8/BF16 Mixed Precision")
|
|
print(f"{'='*60}")
|
|
print(f"[SNR-Quant] Model: {model_name}")
|
|
print(f"[SNR-Quant] Target size: {target_size_GB:.2f} GB")
|
|
print(f"{'='*60}")
|
|
|
|
# Load model to CPU
|
|
print(f"\n[Step 1/6] Loading model to CPU...")
|
|
sd = load_file(model_path, device="cpu")
|
|
print(f" ✓ Loaded {len(sd)} tensors")
|
|
|
|
# Detect input dtype
|
|
sample = next(iter(sd.values()))
|
|
input_dtype = sample.dtype
|
|
print(f" ✓ Input dtype: {input_dtype}")
|
|
|
|
# Analyze layers
|
|
print(f"\n[Step 2/6] Analyzing layers (computing SNR and outliers)...")
|
|
layers, total_params, existing_fp8_count = analyze_model(sd)
|
|
if existing_fp8_count > 0:
|
|
print(f" ✓ Found {len(layers)} quantizable 2D layers ({existing_fp8_count} FP8 layers skipped)")
|
|
else:
|
|
print(f" ✓ Found {len(layers)} 2D weight layers")
|
|
|
|
# Detect fragility
|
|
fragility, snr_range, snr_median = detect_fragility(layers)
|
|
outliers = [l.outlier for l in layers]
|
|
outlier_median = np.median(outliers)
|
|
|
|
print(f"\n[Step 3/6] Detecting model fragility...")
|
|
print(f" SNR range: {snr_range:.2f} dB")
|
|
print(f" SNR median: {snr_median:.2f} dB")
|
|
print(f" Outlier med: {outlier_median:.1f}")
|
|
print(f" → Fragility: {fragility.value.upper()}")
|
|
|
|
# Set adaptive percentile
|
|
if fragility.value == "fragile":
|
|
outlier_percentile = 90
|
|
print(f" → Using 90th percentile (conservative)")
|
|
elif fragility.value == "robust":
|
|
outlier_percentile = 50
|
|
print(f" → Using 50th percentile (aggressive)")
|
|
else:
|
|
outlier_percentile = 75
|
|
print(f" → Using 75th percentile (balanced)")
|
|
|
|
# Calculate sizes
|
|
# layers only contains quantizable 2D weights (excludes existing FP8)
|
|
existing_fp8_numel = sum(t.numel() for t in sd.values() if t.dim() == 2 and t.dtype == torch.float8_e4m3fn)
|
|
existing_fp8_size = existing_fp8_numel # FP8 = 1 byte/param
|
|
|
|
size_2d_bf16 = sum(l.numel * 2 for l in layers) # quantizable 2D layers as BF16
|
|
size_2d_fp8 = sum(l.numel for l in layers) # quantizable 2D layers as FP8
|
|
size_1d_bf16 = sum(t.numel() * 2 for t in sd.values() if t.dim() != 2) # 1D layers always BF16
|
|
|
|
# Current total size (existing FP8 + quantizable as BF16 + 1D)
|
|
current_total_size = existing_fp8_size + size_2d_bf16 + size_1d_bf16
|
|
|
|
# Minimum achievable (all quantizable as FP8)
|
|
min_size = existing_fp8_size + size_2d_fp8 + size_1d_bf16
|
|
|
|
target_bytes = target_size_GB * 1024**3
|
|
need_save = current_total_size - target_bytes
|
|
|
|
print(f"\n[Step 4/6] Calculating size targets...")
|
|
print(f" Current model size: {current_total_size / 1024**3:.2f} GB")
|
|
print(f" └─ Existing FP8: {existing_fp8_size / 1024**3:.2f} GB")
|
|
print(f" └─ Quantizable BF16: {size_2d_bf16 / 1024**3:.2f} GB")
|
|
if existing_fp8_count > 0:
|
|
print(f" Note: {existing_fp8_count} FP8 layers already quantized (skipped)")
|
|
print(f" Minimum achievable: {min_size / 1024**3:.2f} GB")
|
|
print(f" Target size: {target_bytes / 1024**3:.2f} GB")
|
|
print(f" Need to save: {need_save / 1024**3:.2f} GB")
|
|
|
|
# Check if target is achievable
|
|
if target_bytes < min_size:
|
|
print(f"\n ⚠ WARNING: Target {target_size_GB:.2f} GB is below minimum!")
|
|
print(f" Adjusting to {min_size / 1024**3:.2f} GB (all quantizable layers as FP8)")
|
|
target_bytes = min_size
|
|
need_save = current_total_size - min_size
|
|
|
|
if need_save <= 0:
|
|
print(f"\n[Step 5/6] Target >= Current size, no quantization needed")
|
|
print(f" Saving all layers as-is...")
|
|
|
|
# Build layer info lookup
|
|
layer_info_map = {l.name: l for l in layers}
|
|
|
|
quantized = {}
|
|
for name, tensor in sd.items():
|
|
if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
|
|
# Keep existing FP8 layers as is
|
|
quantized[name] = tensor
|
|
elif tensor.dim() == 2:
|
|
quantized[name] = tensor.to(torch.bfloat16)
|
|
else:
|
|
quantized[name] = tensor.to(torch.bfloat16)
|
|
save_file(quantized, output_path)
|
|
|
|
# Generate report - all layers BF16
|
|
report_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB_report.csv")
|
|
print(f"\n Generating quantization report: {report_path}")
|
|
|
|
report_lines = ["layer_name,precision,params_numel,SNR_dB,outlier,precision_change"]
|
|
|
|
# Add existing FP8 layers (already quantized, skipped)
|
|
for name, tensor in sd.items():
|
|
if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
|
|
report_lines.append(f"{name},FP8 (existing),{tensor.numel()},N/A,N/A,already quantized")
|
|
|
|
# Add other layers
|
|
for name, tensor in sd.items():
|
|
if tensor.dim() == 2 and name in layer_info_map:
|
|
l = layer_info_map[name]
|
|
report_lines.append(f"{name},BF16,{tensor.numel()},{l.snr:.2f},{l.outlier:.2f},unchanged")
|
|
elif tensor.dim() != 2:
|
|
report_lines.append(f"{name},BF16,{tensor.numel()},N/A,N/A,unchanged")
|
|
|
|
with open(report_path, 'w', encoding='utf-8') as f:
|
|
f.write('\n'.join(report_lines))
|
|
|
|
print(f" ✓ Report saved: {report_path}")
|
|
|
|
del sd, quantized
|
|
print(f"\n[Step 6/6] Saving to: {output_path}")
|
|
print(f" ✓ Done! Output: quant_{base_name}.safetensors")
|
|
print(f"{'='*60}\n")
|
|
return ()
|
|
|
|
# Iterative selection
|
|
print(f"\n[Step 5/6] Selecting layers for FP8 quantization...")
|
|
final_fp8 = []
|
|
iteration = 0
|
|
|
|
while outlier_percentile >= 50:
|
|
iteration += 1
|
|
threshold = np.percentile(outliers, outlier_percentile)
|
|
quantizable = [l for l in layers if l.outlier < threshold]
|
|
quantizable_sorted = sorted(quantizable, key=lambda x: x.snr, reverse=True)
|
|
max_save = sum(l.numel for l in quantizable)
|
|
|
|
if max_save >= need_save:
|
|
saved = 0
|
|
selected = 0
|
|
for l in quantizable_sorted:
|
|
if saved >= need_save:
|
|
break
|
|
final_fp8.append(l.name)
|
|
saved += l.numel
|
|
selected += 1
|
|
print(f" Iteration {iteration}: {selected} layers selected at {outlier_percentile}th percentile")
|
|
print(f" Saved: {saved / 1024**3:.2f} GB / {need_save / 1024**3:.2f} GB needed")
|
|
break
|
|
else:
|
|
print(f" Iteration {iteration}: Not reachable at {outlier_percentile}th ({max_save / 1024**3:.2f} GB < {need_save / 1024**3:.2f} GB)")
|
|
outlier_percentile -= 5
|
|
|
|
if not final_fp8:
|
|
quantizable_sorted = sorted(layers, key=lambda x: x.snr, reverse=True)
|
|
final_fp8 = [l.name for l in quantizable_sorted]
|
|
print(f" Using ALL {len(final_fp8)} layers for maximum quantization")
|
|
|
|
fp8_set = set(final_fp8)
|
|
bf16_count_layers = len(layers) - len(fp8_set)
|
|
|
|
# Calculate final sizes
|
|
final_2d_fp8 = sum(l.numel for l in layers if l.name in fp8_set)
|
|
final_2d_bf16 = sum(l.numel * 2 for l in layers if l.name not in fp8_set)
|
|
final_size_gb = (existing_fp8_size + final_2d_fp8 + final_2d_bf16 + size_1d_bf16) / 1024**3
|
|
|
|
print(f"\n Final configuration:")
|
|
print(f" FP8 (existing): {existing_fp8_count} ({existing_fp8_size / 1024**3:.2f} GB)")
|
|
print(f" FP8 (new): {len(fp8_set)} ({final_2d_fp8 / 1024**3:.2f} GB)")
|
|
print(f" BF16 (2D): {bf16_count_layers} ({final_2d_bf16 / 1024**3:.2f} GB)")
|
|
print(f" └─ Est. size: {final_size_gb:.2f} GB")
|
|
|
|
# Generate quantization report before quantizing
|
|
report_path = os.path.join(diffusion_models_dir, f"quant_{base_name}_{target_size_GB}GB_report.csv")
|
|
print(f"\n Generating quantization report: {report_path}")
|
|
|
|
# Build layer info lookup
|
|
layer_info_map = {l.name: l for l in layers}
|
|
|
|
report_lines = ["layer_name,precision,params_numel,SNR_dB,outlier,precision_change"]
|
|
|
|
# Add existing FP8 layers (already quantized, skipped)
|
|
for name, tensor in sd.items():
|
|
if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
|
|
report_lines.append(f"{name},FP8 (existing),{tensor.numel()},N/A,N/A,already quantized")
|
|
|
|
# Add quantizable layers
|
|
for l in layers:
|
|
if l.name in fp8_set:
|
|
report_lines.append(f"{l.name},FP8,{l.numel},{l.snr:.2f},{l.outlier:.2f},BF16→FP8")
|
|
else:
|
|
report_lines.append(f"{l.name},BF16,{l.numel},{l.snr:.2f},{l.outlier:.2f},unchanged")
|
|
|
|
# Add 1D layers
|
|
for name, tensor in sd.items():
|
|
if tensor.dim() != 2:
|
|
report_lines.append(f"{name},BF16,{tensor.numel()},N/A,N/A,unchanged")
|
|
|
|
with open(report_path, 'w', encoding='utf-8') as f:
|
|
f.write('\n'.join(report_lines))
|
|
|
|
print(f" ✓ Report saved: {report_path}")
|
|
|
|
# Quantize with progress bar
|
|
print(f"\n[Step 6/6] Quantizing model...")
|
|
quantized = {}
|
|
pbar = ProgressBar(len(sd))
|
|
fp8_count = 0
|
|
bf16_count = 0
|
|
|
|
for i, (name, tensor) in enumerate(sd.items()):
|
|
if tensor.dim() == 2 and tensor.dtype == torch.float8_e4m3fn:
|
|
# Already FP8 - keep as is
|
|
quantized[name] = tensor
|
|
fp8_count += 1
|
|
elif tensor.dim() == 2 and name in fp8_set:
|
|
# 2D layer selected for FP8 quantization
|
|
quantized[name] = tensor.to(torch.float8_e4m3fn)
|
|
fp8_count += 1
|
|
else:
|
|
# BF16: all 1D layers + 2D layers not selected for FP8
|
|
quantized[name] = tensor.to(torch.bfloat16)
|
|
bf16_count += 1
|
|
pbar.update_absolute(i + 1)
|
|
if (i + 1) % 100 == 0 or i == len(sd) - 1:
|
|
print(f" Progress: {i + 1}/{len(sd)} (FP8: {fp8_count}, BF16: {bf16_count})")
|
|
|
|
# Save
|
|
print(f"\n Saving to: {output_path}")
|
|
save_file(quantized, output_path)
|
|
|
|
# Cleanup
|
|
del quantized, sd
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
print(f" ✓ Quantization complete!")
|
|
print(f" Output: quant_{base_name}.safetensors")
|
|
print(f"{'='*60}\n")
|
|
return ()
|
|
|
|
|
|
class SNRAnalyzeNode:
|
|
"""Analyze model and print SNR/outlier statistics."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("diffusion_models"), {
|
|
"tooltip": "Model from ComfyUI/models/diffusion_models"
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("report",)
|
|
FUNCTION = "analyze"
|
|
CATEGORY = "SNR-Quant"
|
|
|
|
def analyze(self, model_name: str):
|
|
model_path = folder_paths.get_full_path("diffusion_models", model_name)
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"[SNR-Quant] Model Analyzer")
|
|
print(f"{'='*60}")
|
|
print(f"[SNR-Quant] Model: {model_name}")
|
|
|
|
print(f"\n[Step 1/3] Loading model...")
|
|
sd = load_file(model_path, device="cpu")
|
|
print(f" ✓ Loaded {len(sd)} tensors")
|
|
|
|
print(f"\n[Step 2/3] Analyzing layers...")
|
|
layers, total_params, existing_fp8_count = analyze_model(sd)
|
|
if existing_fp8_count > 0:
|
|
print(f" ✓ Found {len(layers)} quantizable 2D layers ({existing_fp8_count} FP8 layers skipped)")
|
|
else:
|
|
print(f" ✓ Found {len(layers)} 2D weight layers")
|
|
|
|
print(f"\n[Step 3/3] Computing statistics...")
|
|
fragility, snr_range, snr_median = detect_fragility(layers)
|
|
outliers = [l.outlier for l in layers]
|
|
outlier_median = np.median(outliers)
|
|
snr_min = min([l.snr for l in layers])
|
|
snr_max = max([l.snr for l in layers])
|
|
|
|
print(f" SNR: min={snr_min:.2f}, max={snr_max:.2f}, median={snr_median:.2f} dB")
|
|
print(f" Outliers: min={min(outliers):.2f}, max={max(outliers):.2f}, median={outlier_median:.1f}")
|
|
print(f" Fragility: {fragility.value}")
|
|
|
|
# Build report - calculate actual model sizes
|
|
size_2d_bf16 = sum(l.numel * 2 for l in layers) # quantizable 2D as BF16
|
|
size_2d_fp8 = sum(l.numel for l in layers) # quantizable 2D as FP8
|
|
size_1d_bf16 = sum(t.numel() * 2 for t in sd.values() if t.dim() != 2) # 1D layers BF16
|
|
|
|
# Calculate already quantized FP8 layers size
|
|
existing_fp8_numel = sum(t.numel() for t in sd.values() if t.dim() == 2 and t.dtype == torch.float8_e4m3fn)
|
|
existing_fp8_size = existing_fp8_numel # FP8 = 1 byte/param
|
|
|
|
# Current total size (with existing FP8 + all quantizable as BF16)
|
|
current_total_size = existing_fp8_size + size_2d_bf16 + size_1d_bf16
|
|
|
|
# Minimum achievable (all quantizable as FP8)
|
|
min_size = existing_fp8_size + size_2d_fp8 + size_1d_bf16
|
|
|
|
# Analyze SNR and outlier distribution for recommendations
|
|
snrs = [l.snr for l in layers]
|
|
outliers = [l.outlier for l in layers]
|
|
|
|
# Count layers by SNR quality
|
|
snr_excellent = sum(1 for s in snrs if s >= 35)
|
|
snr_good = sum(1 for s in snrs if 32 <= s < 35)
|
|
snr_moderate = sum(1 for s in snrs if 30 <= s < 32)
|
|
snr_poor = sum(1 for s in snrs if s < 30)
|
|
|
|
# Count layers by outlier risk
|
|
outlier_low = sum(1 for o in outliers if o < 15)
|
|
outlier_medium = sum(1 for o in outliers if 15 <= o < 30)
|
|
outlier_high = sum(1 for o in outliers if o >= 30)
|
|
|
|
# Calculate safety analysis
|
|
safe_to_quantize = sum(1 for l in layers if l.snr >= 32 and l.outlier < 20)
|
|
moderate_risk = sum(1 for l in layers if (30 <= l.snr < 32 or 20 <= l.outlier < 30))
|
|
high_risk = sum(1 for l in layers if l.snr < 30 or l.outlier >= 30)
|
|
|
|
safe_ratio = safe_to_quantize / len(layers) if len(layers) > 0 else 0
|
|
moderate_ratio = moderate_risk / len(layers) if len(layers) > 0 else 0
|
|
high_risk_ratio = high_risk / len(layers) if len(layers) > 0 else 0
|
|
|
|
# Weight by layer size for accurate size estimation
|
|
safe_numel = sum(l.numel for l in layers if l.snr >= 32 and l.outlier < 20)
|
|
moderate_numel = sum(l.numel for l in layers if (30 <= l.snr < 32 or 20 <= l.outlier < 30))
|
|
high_risk_numel = sum(l.numel for l in layers if l.snr < 30 or l.outlier >= 30)
|
|
|
|
# Size calculations (savings from quantizing)
|
|
conservative_fp8_size = safe_numel # Only safe layers as FP8
|
|
balanced_fp8_size = safe_numel + moderate_numel * 0.6 # Safe + 60% of moderate
|
|
aggressive_fp8_size = safe_numel + moderate_numel * 0.9 + high_risk_numel * 0.3 # Most except very risky
|
|
|
|
# Final target sizes = current size - savings
|
|
conservative_target = current_total_size - conservative_fp8_size
|
|
balanced_target = current_total_size - balanced_fp8_size
|
|
aggressive_target = current_total_size - aggressive_fp8_size
|
|
|
|
# Ensure targets are within achievable range
|
|
conservative_target = max(conservative_target, min_size)
|
|
balanced_target = max(balanced_target, min_size)
|
|
aggressive_target = max(aggressive_target, min_size)
|
|
|
|
# Build strategy description
|
|
if len(layers) == 0:
|
|
strategy_note = "No quantizable layers remaining - model is already fully quantized"
|
|
elif high_risk_ratio > 0.3:
|
|
strategy_note = "High-risk layers detected - recommend conservative quantization"
|
|
elif safe_ratio > 0.5:
|
|
strategy_note = "Many safe-to-quantize layers - can use aggressive strategy"
|
|
else:
|
|
strategy_note = "Mixed layer characteristics - balanced approach recommended"
|
|
|
|
# Handle case where all layers are already quantized
|
|
if len(layers) == 0:
|
|
lines = [
|
|
f"Model: {model_name}",
|
|
f"Fragility: N/A (no quantizable layers)",
|
|
f"2D Layers: 0 + {existing_fp8_count} FP8 (already quantized)",
|
|
f"1D Layers: {len(sd) - existing_fp8_count}",
|
|
f"",
|
|
f"This model is already fully quantized to FP8.",
|
|
f"Current size: {current_total_size / 1024**3:.2f} GB",
|
|
f"No further quantization recommended.",
|
|
]
|
|
report = "\n".join(lines)
|
|
print(f"\n{report}")
|
|
print(f"\n{'='*60}\n")
|
|
del sd
|
|
return (report,)
|
|
|
|
lines = [
|
|
f"Model: {model_name}",
|
|
f"Fragility: {fragility.value}",
|
|
f"2D Layers: {len(layers)} + {existing_fp8_count} FP8 (skipped)",
|
|
f"1D Layers: {len(sd) - len(layers) - existing_fp8_count}",
|
|
f"",
|
|
f"Note: {existing_fp8_count} FP8 layers already quantized (skipped in analysis)" if existing_fp8_count > 0 else "",
|
|
f"",
|
|
f"SNR Distribution:",
|
|
f" Min: {snr_min:.2f} dB",
|
|
f" Max: {snr_max:.2f} dB",
|
|
f" Median: {snr_median:.2f} dB",
|
|
f" Range: {snr_range:.2f} dB",
|
|
f"",
|
|
f"SNR Quality Breakdown:",
|
|
f" Excellent (>=35 dB): {snr_excellent} layers ({snr_excellent/max(len(layers),1)*100:.0f}%)",
|
|
f" Good (32-35 dB): {snr_good} layers ({snr_good/max(len(layers),1)*100:.0f}%)",
|
|
f" Moderate (30-32 dB): {snr_moderate} layers ({snr_moderate/max(len(layers),1)*100:.0f}%)",
|
|
f" Poor (<30 dB): {snr_poor} layers ({snr_poor/max(len(layers),1)*100:.0f}%)",
|
|
f"",
|
|
f"Outlier Distribution:",
|
|
f" Min: {min(outliers):.2f}",
|
|
f" Max: {max(outliers):.2f}",
|
|
f" Median: {outlier_median:.1f}",
|
|
f"",
|
|
f"Outlier Risk Breakdown:",
|
|
f" Low (<15): {outlier_low} layers ({outlier_low/max(len(layers),1)*100:.0f}%)",
|
|
f" Medium (15-30): {outlier_medium} layers ({outlier_medium/max(len(layers),1)*100:.0f}%)",
|
|
f" High (>=30): {outlier_high} layers ({outlier_high/max(len(layers),1)*100:.0f}%)",
|
|
f"",
|
|
f"Quantization Safety Analysis:",
|
|
f" Safe to quantize (SNR>=32 AND outlier<20): {safe_to_quantize} layers ({safe_ratio*100:.0f}%)",
|
|
f" Moderate risk: {moderate_risk} layers ({moderate_ratio*100:.0f}%)",
|
|
f" High risk (protect recommended): {high_risk} layers ({high_risk_ratio*100:.0f}%)",
|
|
f"",
|
|
f"Size estimates:",
|
|
f" Current model size: {current_total_size / 1024**3:.2f} GB",
|
|
f" └─ Existing FP8: {existing_fp8_size / 1024**3:.2f} GB",
|
|
f" └─ Quantizable BF16: {size_2d_bf16 / 1024**3:.2f} GB",
|
|
f" Minimum achievable: {min_size / 1024**3:.2f} GB (all quantizable as FP8)",
|
|
f"",
|
|
f"Quantization Recommendations:",
|
|
f" Strategy: {strategy_note}",
|
|
f"",
|
|
f" Recommended target sizes (final model size):",
|
|
f" Conservative: {conservative_target / 1024**3:.2f} GB (save {(current_total_size - conservative_target) / 1024**3:.2f} GB)",
|
|
f" └─ FP8 layers: only safe layers (SNR>=32, outlier<20)",
|
|
f" └─ Expected quality: Near-lossless",
|
|
f" └─ Best for: Production, quality-critical work",
|
|
f"",
|
|
f" Balanced: {balanced_target / 1024**3:.2f} GB (save {(current_total_size - balanced_target) / 1024**3:.2f} GB)",
|
|
f" └─ FP8 layers: safe + partial moderate-risk layers",
|
|
f" └─ Expected quality: Minimal degradation",
|
|
f" └─ Best for: General use, good quality/size tradeoff",
|
|
f"",
|
|
f" Aggressive: {aggressive_target / 1024**3:.2f} GB (save {(current_total_size - aggressive_target) / 1024**3:.2f} GB)",
|
|
f" └─ FP8 layers: all except high-risk",
|
|
f" └─ Expected quality: Acceptable for most cases",
|
|
f" └─ Best for: VRAM-constrained environments",
|
|
]
|
|
report = "\n".join(lines)
|
|
|
|
print(f"\n{report}")
|
|
print(f"\n{'='*60}\n")
|
|
del sd
|
|
|
|
return (report,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SNRQuant": SNRQuantNode,
|
|
"SNRAnalyze": SNRAnalyzeNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SNRQuant": "SNR Quant (FP8/ BF16)",
|
|
"SNRAnalyze": "SNR Model Analyzer",
|
|
}
|
|
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|