Terpentinas-EasyLoRAMerger/engine/fp8_quantizer.py

"""
FP8 Quantizer — Shared Module for Scale-Then-Cast Quantization

Provides the canonical FP8 quantization algorithm used across EasyLoRAMerger:
  - compute_weight_scale()   — Per-tensor scale factor for FP8 quantization
  - quantize_to_fp8()        — Scale-then-cast quantization (matching Creator's algorithm)
  - should_preserve_bf16()   — BF16 preservation for sensitive layer types
  - get_dequant_dtype()      — GPU-aware dequant target resolution (FP16 for legacy GPUs)
  - FP8_PRESERVE_BF16_PATTERNS — The pattern tuple (source of truth)
  - build_fp8_quantization_metadata() — Inject _quantization_metadata into metadata dict
  - strip_input_scale_keys()         — Remove .input_scale companion keys for ComfyUI compat

All stateless functions operating on tensors and keys.
This module is the single source of truth; musubi_checkpoint_studio.py
and baker_node.py delegate to it.
"""

import json
from typing import Dict, Optional

import torch

# ── FP8 format constants ──────────────────────────────────────────────────────
# Max representable values for each FP8 format (used in scale computation).
#   float8_e4m3fn:  max = 448.0  (3 mantissa bits, 4 exponent bits)
#   float8_e5m2:    max = 57344.0 (2 mantissa bits, 5 exponent bits)
FP8_MAX_VALUES = {
    torch.float8_e4m3fn: 448.0,
    torch.float8_e5m2: 57344.0,
}


# 🔥 FP8 IMPROVEMENT: Patterns that should be preserved in BF16 instead of quantized to FP8
#
# Creator's FP8 selectively keeps critical layers in BF16:
#   - Input projections: img_in, time_in, txt_in, vector_in, guidance_in
#   - Output projections: final_layer
#   - Modulation layers: single_stream_modulation, double_stream_modulation_*
#   - Norm scale/bias tensors
#
# Patterns support two key formats:
#   - Diffusers-style: uses '/' path separators (model/diffusion_model/...)
#   - FLUX-style: uses '.' dot notation (double_blocks.0.img_attn...)
#
# SDXL/SD1.5 and Anima patterns were validated in INT8_PRESERVE_FP16_PATTERNS
# (engine/int8_quantizer.py) — ported here for equivalent FP8 protection.
FP8_PRESERVE_BF16_PATTERNS: tuple = (
    # ── Diffusers-style path separators ──
    '/final_layer/',
    '/time_in/',
    '/txt_in/',
    '/vector_in/',
    '/guidance_in/',
    '/norm.',
    # ── 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 diffusers-style path equivalents ──
    '/time_embed/',
    '/label_emb/',
    '/input_blocks/0/',
    '/middle_block/',
    '/out/',

    # ── 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',
)


# ── Public API ────────────────────────────────────────────────────────────────


def get_dequant_dtype(
    target_device: Optional[torch.device] = None,
    file_metadata: Optional[Dict] = None,
) -> torch.dtype:
    """Return optimal dequantization dtype for FP8 tensors on the target device.

    Resolution order:
    1. If ``file_metadata`` has a ``'dequant_target'`` hint → use that
       (honors the baker's intent for cross‑compilation scenarios).
    2. Else → auto‑detect from device capability:
       - CC >= 8.0 (Ampere+): ``bfloat16`` (natively supported)
       - CC < 8.0 (Pascal/Turing): ``float16`` (only natively‑supported 16‑bit option)
       - CPU: ``float32``

    Args:
        target_device: The torch device that will run inference.
            If ``None``, uses ``DeviceManager.get_device()``.
        file_metadata: Optional metadata dict from the safetensors file.
            If present and contains ``'dequant_target'``, that value takes
            priority over auto‑detection.

    Returns:
        torch.dtype: ``torch.bfloat16``, ``torch.float16``, or ``torch.float32``.
    """
    # Layer 1: File metadata hint (cross‑compilation support)
    if file_metadata and 'dequant_target' in file_metadata:
        hint = file_metadata['dequant_target']
        if hint == 'float16':
            return torch.float16
        elif hint == 'bfloat16':
            return torch.bfloat16
        elif hint == 'float32':
            return torch.float32
        # Unknown hint → fall through to auto‑detect

    # Layer 2: Auto‑detect from device capability
    if target_device is None:
        # Lazy import to avoid circular dependency
        from ..utils import DeviceManager as _DevMgr
        target_device = _DevMgr.get_device()
    elif isinstance(target_device, str):
        from ..utils import DeviceManager as _DevMgr
        target_device = _DevMgr.get_device(target_device)

    if target_device.type == 'cuda':
        try:
            cap = torch.cuda.get_device_capability(target_device)
            # CC >= 8.0 (Ampere, Ada, Hopper, Blackwell) → BF16 native
            if cap[0] >= 8:
                return torch.bfloat16
            # CC < 8.0 (Pascal, Turing, Volta) → FP16 native or emulated
            return torch.float16
        except Exception:
            return torch.float16  # safe fallback

    return torch.float32


def should_preserve_bf16(key: str) -> bool:
    """Return True if this key should stay in BF16 when target precision is FP8.

    Per FP8 comparison analysis (flux-test/FP8_COMPARISON_REPORT.md), the Creator's FP8
    preserves these layers in BF16 because they are sensitive small tensors (norm scales,
    input projections) where FP8 quantization introduces measurable precision loss.
    """
    key_lower = key.lower()
    for pattern in FP8_PRESERVE_BF16_PATTERNS:
        if pattern in key_lower:
            return True
    return False


def compute_weight_scale(tensor: torch.Tensor, target_dtype: torch.dtype) -> torch.Tensor:
    """Compute the per-tensor weight_scale for FP8 quantization.

    The scale factor maps the max absolute value of the tensor to the
    max representable value of the FP8 format:
        weight_scale = max(abs(tensor)) / max_fp8

    Args:
        tensor: Original float tensor (before FP8 quantization).
        target_dtype: torch.float8_e4m3fn or torch.float8_e5m2.

    Returns:
        F32 scalar tensor containing the scale factor.
    """
    max_fp8 = FP8_MAX_VALUES.get(target_dtype, 448.0)
    # Convert to float32 first — abs() and max() are not implemented for
    # Float8 dtypes on all PyTorch versions (e.g. ComfyUI's CUDA build).
    tensor_f32 = tensor.float()
    # Guard: sanitize NaN/Inf → 0 to prevent corrupted scale factors.
    # A single NaN in max_abs → scale=NaN → tensor/NaN = NaN during dequant.
    if not torch.isfinite(tensor_f32).all():
        tensor_f32 = torch.nan_to_num(tensor_f32, nan=0.0, posinf=0.0, neginf=0.0)
    max_abs = tensor_f32.abs().max()
    # Force F32 tensor division to match Creator's precision (closes ~2% gap)
    scale = max_abs / torch.tensor(max_fp8, dtype=torch.float32, device=max_abs.device)
    scale = torch.clamp(scale, min=1e-12)
    return scale.to(dtype=torch.float32)


def quantize_to_fp8(tensor: torch.Tensor, fp8_dtype: torch.dtype) -> torch.Tensor:
    """Quantize tensor to FP8 using scale-then-cast (matching Creator's algorithm).

    Creator's approach (validated at 100.00% match):
      1. Compute scale = max_abs / max_fp8 (in F32 precision)
      2. Divide tensor by scale
      3. Clamp to [-max_fp8, max_fp8]
      4. Cast to FP8

    This differs from PyTorch's .to(fp8) which uses hardware rounding directly
    without scaling — producing ~98% different byte values.

    Edge case guards:
      - NaN/Inf in weights → sanitized to 0.0 before scale computation.
        (Without this, a single NaN would produce scale=NaN, corrupting the
        entire block during dequant via `fp8_value * NaN = NaN`.)
      - All-zero tensors → scale clamped to 1e-12, quantized to 0. Correct.
    """
    max_fp8 = FP8_MAX_VALUES.get(fp8_dtype, 448.0)
    tensor_f32 = tensor.float()

    # Guard: sanitize NaN/Inf → 0 to prevent corrupted scale factors.
    if not torch.isfinite(tensor_f32).all():
        tensor_f32 = torch.nan_to_num(tensor_f32, nan=0.0, posinf=0.0, neginf=0.0)

    # Step 1: Compute scale from cleaned tensor (F32 precision)
    max_abs = tensor_f32.abs().max()
    scale = max_abs / torch.tensor(max_fp8, dtype=torch.float32, device=max_abs.device)
    scale = torch.clamp(scale, min=1e-12)  # All-zero tensor → 1e-12, not 0

    # Steps 2-4: scale-then-cast
    tensor_f32.div_(scale)          # Normalize to [-max_fp8, max_fp8] range
    tensor_f32.clamp_(-max_fp8, max_fp8)  # Clamp to FP8 representable range
    result = tensor_f32.to(dtype=fp8_dtype)

    return result


def quantize_weight_to_fp8_with_scales(
    tensor: torch.Tensor,
    fp8_dtype: torch.dtype,
) -> tuple:
    """Quantize a weight tensor to FP8 and return (quantized_tensor, weight_scale, input_scale).

    This is a convenience wrapper that performs all three steps together:
      1. Compute weight_scale
      2. Compute input_scale (same as weight_scale per Creator's convention)
      3. Quantize the tensor

    Args:
        tensor: Original float tensor (before FP8 quantization).
        fp8_dtype: torch.float8_e4m3fn or torch.float8_e5m2.

    Returns:
        Tuple of (quantized_fp8_tensor, weight_scale_f32_scalar, input_scale_f32_scalar).
    """
    wscale = compute_weight_scale(tensor, fp8_dtype)
    iscale = compute_weight_scale(tensor, fp8_dtype)  # Same computation per Creator
    quantized = quantize_to_fp8(tensor, fp8_dtype)
    return quantized, wscale, iscale


def quantize_weight_to_fp8_with_scales_optimized(
    tensor: torch.Tensor,
    fp8_dtype: torch.dtype,
) -> tuple:
    """Single-pass FP8 quantization — computes weight_scale and quantize in ONE scan.

    Eliminates the 2 redundant tensor scans from the original implementation
    (which called compute_weight_scale twice + quantize_to_fp8 separately).

    Args:
        tensor: Original float tensor (before FP8 quantization).
        fp8_dtype: torch.float8_e4m3fn or torch.float8_e5m2.

    Returns:
        Tuple of (quantized_fp8_tensor, weight_scale_f32_scalar, input_scale_f32_scalar).
    """
    max_fp8 = FP8_MAX_VALUES.get(fp8_dtype, 448.0)

    # ── Single scan: float → abs().max() → scale ──
    tensor_f32 = tensor.float()
    if not torch.isfinite(tensor_f32).all():
        tensor_f32 = torch.nan_to_num(tensor_f32, nan=0.0, posinf=0.0, neginf=0.0)

    max_abs = tensor_f32.abs().max()
    scale = max_abs / torch.tensor(max_fp8, dtype=torch.float32, device=max_abs.device)
    scale = torch.clamp(scale, min=1e-12)

    # ── Scale-then-cast (in-place on the F32 copy) ──
    tensor_f32.div_(scale)
    tensor_f32.clamp_(-max_fp8, max_fp8)
    quantized = tensor_f32.to(dtype=fp8_dtype)

    # weight_scale and input_scale are identical per Creator's convention
    wscale = scale.to(dtype=torch.float32)
    iscale = scale.to(dtype=torch.float32)  # same value, separate tensor

    return quantized, wscale, iscale


def dequant_fp8_tensor(
    tensor: torch.Tensor,
    key: str,
    target_dtype: torch.dtype,
    scale_store: Dict[str, torch.Tensor],
) -> torch.Tensor:
    """Dequantize an FP8 tensor to target dtype using companion scale factors.

    FP8 tensors store per-channel quantization scales as companion tensors
    (e.g. '.weight_scale' and '.input_scale' suffixes). This function:
      1. Casts FP8 → target dtype (e.g. bf16) — gives scaled values
      2. Looks up the per-channel scale factor from scale_store
      3. Multiplies: out = fp8_as_target * scale (broadcasted)

    Args:
        tensor: FP8 tensor (float8_e4m3fn or float8_e5m2).
        key: Checkpoint key for the tensor (used to derive scale key names).
        target_dtype: Target dtype for dequantization (e.g. bfloat16).
        scale_store: Dict-like containing companion scale tensors.
            Supports plain dict and _LazyCheckpointMapping (both have __contains__).

    Returns:
        Dequantized tensor at target_dtype. Returns tensor unchanged if no
        companion scale is found (fallback for non-quantized keys).
    """
    tensor = tensor.to(dtype=target_dtype)
    # Derive the base prefix for scale key lookup
    base = key
    if key.endswith('.weight'):
        base = key[:-len('.weight')]
    # Look for companion scale (weight_scale or input_scale)
    for suffix in ('.weight_scale', '.input_scale'):
        scale_key = base + suffix
        if scale_key in scale_store:
            scale = scale_store[scale_key].to(dtype=target_dtype)
            # Broadcast: weight [C, K1, K2, ...] * scale [C, 1, 1, ...]
            for _ in range(tensor.dim() - scale.dim()):
                scale = scale.unsqueeze(-1)
            tensor = tensor * scale
            break
    return tensor


# ══════════════════════════════════════════════════════════════════════════════
# Shared FP8 metadata helpers — used by both baker_node.py and checkpoint_merger_node.py
# Extracted here to follow "One Pattern Per Concept" (working_principles.md).
# ══════════════════════════════════════════════════════════════════════════════


def build_fp8_quantization_metadata(
    output_sd,
    metadata: dict,
    label: str = "",
) -> dict:
    """Build _quantization_metadata for MixedPrecisionOps and inject into metadata.

    Handles both _LazyCheckpointMapping (scans header/write_cache for dtype)
    and plain dicts (scans tensor values directly).

    Returns the (possibly updated) metadata dict.
    """
    fp8_dtype_str = getattr(output_sd, '_user_fp8_dtype_str', 'F8_E4M3')
    fp8_format = {'F8_E4M3': 'float8_e4m3fn', 'F8_E5M2': 'float8_e5m2'}.get(
        fp8_dtype_str, 'float8_e4m3fn')

    fp8_layers = {}

    if hasattr(output_sd, '_iter_keys_unfiltered'):
        # _LazyCheckpointMapping: check write cache and file header
        for k in output_sd._iter_keys_unfiltered():
            if k.endswith(('.weight_scale', '.input_scale')):
                continue
            if output_sd._write_cache and k in output_sd._write_cache:
                dtype = output_sd._write_cache[k].dtype
            else:
                output_sd._ensure_open()
                info = output_sd._header.get(k, {})
                dtype_str = info.get('dtype', 'BF16')
                dtype = output_sd._SAFETENSORS_DTYPE_MAP.get(dtype_str, torch.bfloat16)
            is_fp8 = dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
            if is_fp8:
                layer_name = k[:-len('.weight')] if k.endswith('.weight') else k
                fp8_layers[layer_name] = {'format': fp8_format}
    else:
        # Plain dict: check tensor dtype directly
        for k, v in output_sd.items():
            if isinstance(v, torch.Tensor) and v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
                if not k.endswith(('.weight_scale', '.input_scale')):
                    layer_name = k[:-len('.weight')] if k.endswith('.weight') else k
                    # Derive format from actual tensor dtype — more reliable than
                    # _user_fp8_dtype_str (which can't be set on plain dicts).
                    _layer_fmt = 'float8_e4m3fn' if v.dtype == torch.float8_e4m3fn else 'float8_e5m2'
                    fp8_layers[layer_name] = {'format': _layer_fmt}

    if fp8_layers:
        qmeta = {'format_version': '1.0', 'layers': fp8_layers}
        metadata = dict(metadata) if metadata else {}
        metadata['_quantization_metadata'] = json.dumps(qmeta)
        suffix = f" ({label})" if label else ""
        print(f"   📋 Added quantization metadata ({len(fp8_layers)} FP8 layers){suffix}")

    return metadata


def strip_input_scale_keys(output_sd) -> int:
    """Strip .input_scale keys (MixedPrecisionOps only needs .weight_scale).

    Returns the number of keys stripped.
    """
    input_scale_keys = [k for k in output_sd if k.endswith('.input_scale')]
    for k in input_scale_keys:
        output_sd.pop(k, None)
    if input_scale_keys:
        print(f"   🧹 Stripped {len(input_scale_keys)} .input_scale tensors for MixedPrecisionOps")
    return len(input_scale_keys)