vllm-omni/vllm_omni/quantization/tools/merge_mxfp4_dualscale_checkpoint.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Merge W4A4_MXFP4_DUALSCALE quantized Wan2.2 weights into HF Diffusers format.

msModelSlim produces a checkpoint where each linear layer is either:
  - W4A4_MXFP4_DUALSCALE: quantized weights with fine/coarse scales and mul_scale
  - FLOAT (BF16 fallback): kept as original BF16 weights for precision-sensitive layers

BF16 fallback layers may be interleaved anywhere in the transformer (not just leading
blocks). The merge script detects them from quant_model_description.json and writes
their prefixes into config.json as ignored_layers so the runtime dynamically routes
each layer to the correct quantization method at weight-loading time.

MXFP4_DUALSCALE key structure per linear layer (msModelSlim naming)
---------------------------------------------------------------------
  blocks.N.X.linear.weight             W4A4_MXFP4_DUALSCALE  – int8 (FP4 packed)
  blocks.N.X.linear.weight_scale       W4A4_MXFP4_DUALSCALE  – uint8 (float8_e8m0fnu fine scale, per-32K)
  blocks.N.X.linear.weight_dual_scale  W4A4_MXFP4_DUALSCALE  – float32 (coarse scale, per-512K)
  blocks.N.X.linear.bias               FLOAT                  – bias (if present)
  blocks.N.X.div.mul_scale             FLOAT                  – float32 per-input-channel activation pre-scale

BF16 fallback key structure (no wrappers, plain linear weights):
  blocks.N.X.weight                    FLOAT
  condition_embedder.*.weight          FLOAT  (always BF16 — not quantized)

Self-attention QKV notes
------------------------
Self-attention Q/K/V weights are separate in the checkpoint (self_attn.q/k/v) but fused
into a single to_qkv layer in vllm-omni. The transformer's load_weights() handles this via
stacked_params_mapping. This script keeps Q/K/V keys separate — do NOT pre-fuse them here.

For mul_scale specifically: even though Q/K/V process the same input (same mul_scale value),
they are kept separate. load_weights() routes all three to the same to_qkv.mul_scale parameter
and each overwrites the previous. Since Q=K=V for mul_scale, the final value is correct.

NOTE: Pre-fusing as to_qkv.mul_scale would BREAK loading because ".attn1.to_q" is a
substring of ".attn1.to_qkv", causing load_weights() stacked_params_mapping to produce a
garbage key ("to_qkvkv") that is not in params_dict, triggering a break that skips the
direct-load else branch entirely.

Supported model types:
  - Wan2.2-T2V-A14B  (MoE cascade: transformer + transformer_2)
  - Wan2.2-I2V-A14B  (MoE cascade: transformer + transformer_2)

Note: Wan2.2-TI2V-5B is NOT supported. Its smaller parameter count causes
unacceptable accuracy loss under W4A4 quantization. Use MXFP8 for TI2V-5B.

Usage:
  python merge_mxfp4_dualscale_checkpoint.py \\
      --model-type        Wan2.2-T2V-A14B \\
      --original-model    /path/to/Wan2.2-T2V-A14B-Diffusers \\
      --quant-path        /path/to/msmodelslim-output \\
      --output-path       /path/to/merged-output
"""

from __future__ import annotations

import argparse
import json
import pathlib
import re
import shutil
from typing import Any

import torch
from safetensors.torch import load_file, save_file

# ---------------------------------------------------------------------------
# Key rename: msModelSlim naming → Diffusers / vllm-omni naming
# Identical to merge_mxfp8_checkpoint.py; applied to all blocks uniformly.
# ---------------------------------------------------------------------------

TRANSFORMER_KEYS_RENAME_DICT: dict[str, str] = {
    "time_embedding.0": "condition_embedder.time_embedder.linear_1",
    "time_embedding.2": "condition_embedder.time_embedder.linear_2",
    "text_embedding.0": "condition_embedder.text_embedder.linear_1",
    "text_embedding.2": "condition_embedder.text_embedder.linear_2",
    "time_projection.1": "condition_embedder.time_proj",
    "head.modulation": "scale_shift_table",
    "head.head": "proj_out",
    "modulation": "scale_shift_table",
    "ffn.0": "ffn.net.0.proj",
    "ffn.2": "ffn.net.2",
    # Norm order swap (quant tool: norm1, norm3, norm2 → diffusers: norm1, norm2, norm3)
    "norm2": "norm__placeholder",
    "norm3": "norm2",
    "norm__placeholder": "norm3",
    # Self-attention
    "self_attn.q": "attn1.to_q",
    "self_attn.k": "attn1.to_k",
    "self_attn.v": "attn1.to_v",
    "self_attn.o": "attn1.to_out.0",
    "self_attn.norm_q": "attn1.norm_q",
    "self_attn.norm_k": "attn1.norm_k",
    # Cross-attention
    "cross_attn.q": "attn2.to_q",
    "cross_attn.k": "attn2.to_k",
    "cross_attn.v": "attn2.to_v",
    "cross_attn.o": "attn2.to_out.0",
    "cross_attn.norm_q": "attn2.norm_q",
    "cross_attn.norm_k": "attn2.norm_k",
    "attn2.to_k_img": "attn2.add_k_proj",
    "attn2.to_v_img": "attn2.add_v_proj",
    "attn2.norm_k_img": "attn2.norm_added_k",
    # I2V image embedder
    "img_emb.proj.0": "condition_embedder.image_embedder.norm1",
    "img_emb.proj.1": "condition_embedder.image_embedder.ff.net.0.proj",
    "img_emb.proj.3": "condition_embedder.image_embedder.ff.net.2",
    "img_emb.proj.4": "condition_embedder.image_embedder.norm2",
    "img_emb.emb_pos": "condition_embedder.image_embedder.pos_embed",
}

SUPPORTED_MODEL_TYPES = ["Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"]
CASCADE_MODEL_TYPES = {"Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"}

# Suffixes that appear inside .linear.* wrapper for MXFP4 tensors.
_LINEAR_ATTRS = ("weight_dual_scale", "weight_scale", "weight", "bias")

_BLOCK_IDX_RE = re.compile(r"^blocks\.(\d+)\.")


# ---------------------------------------------------------------------------
# Key transformation helpers
# ---------------------------------------------------------------------------


def _parse_block_idx(key: str) -> int | None:
    """Extract block index from a key like 'blocks.5.attn1.to_q.weight'."""
    m = _BLOCK_IDX_RE.match(key)
    return int(m.group(1)) if m else None


def _apply_rename_dict(key: str) -> str:
    for src, dst in TRANSFORMER_KEYS_RENAME_DICT.items():
        key = key.replace(src, dst)
    return key


def _strip_mxfp4_wrapper(key: str) -> str:
    """Strip .linear.ATTR or .div.mul_scale wrappers added by msModelSlim.

    MXFP4 tensors are wrapped in sub-modules:
      X.linear.weight / weight_scale / weight_dual_scale / bias
      X.div.mul_scale

    MXFP8 and FLOAT tensors have no wrappers — this function is a no-op for them.

    Examples after apply_rename_dict:
      attn1.to_q.linear.weight      → attn1.to_q.weight
      attn1.to_q.linear.weight_scale → attn1.to_q.weight_scale
      attn1.to_q.div.mul_scale      → attn1.to_q.mul_scale
      attn1.norm_q.weight           → attn1.norm_q.weight  (unchanged)
    """
    # Check longest attribute names first to avoid partial suffix matches.
    for attr in _LINEAR_ATTRS:
        suffix = f".linear.{attr}"
        if key.endswith(suffix):
            return key[: -len(suffix)] + f".{attr}"
    if key.endswith(".div.mul_scale"):
        return key[: -len(".div.mul_scale")] + ".mul_scale"
    return key


# ---------------------------------------------------------------------------
# Quantization metadata helpers
# ---------------------------------------------------------------------------

# Known quantized weight types (FLOAT tensors don't determine block type).
_MXFP8_TYPE = "mxfp8"
_MXFP4_DUALSCALE_TYPE = "mxfp4_dualscale"


def _classify_blocks(quant_meta: dict[str, str]) -> dict[int, str]:
    """Classify each transformer block by quantization type from quant_meta.

    Returns a dict mapping block_idx → 'mxfp8' | 'mxfp4_dualscale'.
    A block's type is determined by the first quantized (non-FLOAT) tensor found for it.
    """
    block_types: dict[int, str] = {}
    for key, qtype in quant_meta.items():
        idx = _parse_block_idx(key)
        if idx is None or idx in block_types:
            continue
        if qtype.startswith("W8A8_MXFP8"):
            block_types[idx] = _MXFP8_TYPE
        elif qtype.startswith("W4A4_MXFP4_DUALSCALE"):
            block_types[idx] = _MXFP4_DUALSCALE_TYPE
    return block_types


def _print_block_summary(block_types: dict[int, str]) -> None:
    """Print a compact run-length summary of the block layout."""
    if not block_types:
        print("  Block layout: (empty)")
        return

    sorted_indices = sorted(block_types)
    runs: list[tuple[int, int, str]] = []
    run_start = sorted_indices[0]
    run_type = block_types[run_start]
    for idx in sorted_indices[1:]:
        if block_types[idx] != run_type:
            runs.append((run_start, idx - 1, run_type))
            run_start = idx
            run_type = block_types[idx]
    runs.append((run_start, sorted_indices[-1], run_type))

    print(f"  Block layout ({len(sorted_indices)} blocks classified):")
    for start, end, btype in runs:
        count = end - start + 1
        range_str = f"{start}" if start == end else f"{start}–{end}"
        print(f"    blocks {range_str:>8}: {btype}  ({count} block{'s' if count > 1 else ''})")


# ---------------------------------------------------------------------------
# Safetensors I/O
# ---------------------------------------------------------------------------


def _load_safetensors_dir(directory: pathlib.Path, glob: str = "*.safetensors") -> dict[str, torch.Tensor]:
    candidates = sorted(directory.glob(glob))
    if not candidates:
        raise FileNotFoundError(f"No safetensors matching '{glob}' found in {directory}")
    state: dict[str, torch.Tensor] = {}
    for f in candidates:
        state.update(load_file(str(f)))
    return state


def _load_quant_safetensors(directory: pathlib.Path) -> dict[str, torch.Tensor]:
    try:
        return _load_safetensors_dir(directory, "quant_model_weight*.safetensors")
    except FileNotFoundError:
        return _load_safetensors_dir(directory)


def _load_quant_meta(directory: pathlib.Path) -> dict[str, str]:
    candidates = sorted(directory.glob("quant_model_description*.json"))
    if not candidates:
        print(f"  WARNING: No quant_model_description*.json in {directory}; treating all tensors as FLOAT.")
        return {}
    with open(candidates[0]) as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# Per-transformer conversion
# ---------------------------------------------------------------------------


def _is_mxfp4_tensor(key: str, qtype: str) -> bool:
    """Return True for MXFP4 quantized tensors and their companion tensors.

    MXFP4 DualScale layers produce four tensor types:
      .linear.weight / .linear.weight_scale / .linear.weight_dual_scale (W4A4_MXFP4_DUALSCALE)
      .linear.bias / .div.mul_scale (FLOAT but companion to a quantized layer)

    BF16 fallback layers have plain .weight / .bias without .linear. / .div. wrappers.
    """
    if qtype.startswith("W4A4_MXFP4_DUALSCALE"):
        return True
    # Companion tensors: bias and mul_scale that belong to an MXFP4 layer.
    return ".linear." in key or ".div.mul_scale" in key


def _diffusers_to_vllm_ignored(diffusers_ignored: list[str]) -> list[str]:
    """Translate diffusers checkpoint prefixes to vllm-omni model parameter names.

    Three transformations align the naming conventions:

    1. Self-attention QKV fusion (attn1 only, not attn2):
       attn1.to_q + attn1.to_k + attn1.to_v → attn1.to_qkv
       Cross-attention (attn2) to_q/k/v are not fused and remain separate.

    2. FFN naming:
       ffn.net.0.proj  → ffn.net_0.proj
       ffn.net.2       → ffn.net_2

    3. to_out index:
       attn1.to_out.0 / attn2.to_out.0  → attn1.to_out / attn2.to_out
    """
    ignored_set = set(diffusers_ignored)
    result: set[str] = set()

    for name in diffusers_ignored:
        m = re.match(r"^(.*\.attn1)\.to_([qkv])$", name)
        if m:
            prefix = m.group(1)
            if all(f"{prefix}.to_{c}" in ignored_set for c in ("q", "k", "v")):
                result.add(f"{prefix}.to_qkv")
            else:
                present = [f"to_{c}" for c in ("q", "k", "v") if f"{prefix}.to_{c}" in ignored_set]
                missing = [f"to_{c}" for c in ("q", "k", "v") if f"{prefix}.to_{c}" not in ignored_set]
                raise ValueError(
                    f"Partial BF16 fallback for '{prefix}': "
                    f"{', '.join(present)} in ignored_layers but {', '.join(missing)} is not. "
                    f"Self-attention Q/K/V are fused into a single to_qkv layer at runtime; "
                    f"all three must share the same precision. "
                    f"Either quantize all of to_q/to_k/to_v or keep all three in BF16."
                )
            continue

        name = re.sub(r"\.ffn\.net\.0\.proj$", ".ffn.net_0.proj", name)
        name = re.sub(r"\.ffn\.net\.2$", ".ffn.net_2", name)
        name = re.sub(r"\.to_out\.0$", ".to_out", name)
        result.add(name)

    return sorted(result)


def _collect_ignored_layers(merged: dict[str, Any], mxfp4_layer_prefixes: set[str]) -> list[str]:
    """Collect vllm-omni parameter name prefixes for BF16 fallback layers.

    A layer is BF16 if it has a .weight tensor but its prefix is not in
    mxfp4_layer_prefixes. Prefixes are returned in vllm-omni parameter naming
    (QKV-fused, FFN underscored, to_out unindexed) so the list can be written
    directly into ignored_layers in config.json without further translation.
    """
    all_weight_prefixes = {key[: -len(".weight")] for key in merged if key.endswith(".weight")}
    return _diffusers_to_vllm_ignored(sorted(all_weight_prefixes - mxfp4_layer_prefixes))


def _convert_transformer(
    quant_subdir: pathlib.Path,
    output_dir: pathlib.Path,
    original_transformer_dir: pathlib.Path,
) -> None:
    """Convert one transformer directory to the mxfp4_dualscale + BF16 mixed format."""
    output_dir.mkdir(parents=True, exist_ok=True)

    # BF16 base: provides the scaffold for non-MXFP4 tensors (norms, embeddings,
    # BF16 fallback linear layers that msModelSlim may omit or keep as FLOAT).
    print(f"  Loading BF16 base from {original_transformer_dir} …")
    base_state = _load_safetensors_dir(original_transformer_dir)
    print(f"  {len(base_state)} BF16 tensors loaded")

    print(f"  Loading quantized weights from {quant_subdir} …")
    quant_state = _load_quant_safetensors(quant_subdir)
    quant_meta = _load_quant_meta(quant_subdir)
    print(f"  {len(quant_state)} quant tensors, {len(quant_meta)} meta entries")

    # Classify blocks for a compact summary (informational only in the new scheme).
    block_types = _classify_blocks(quant_meta)
    _print_block_summary(block_types)

    # Remap MXFP4 quantized tensors only.
    # BF16 fallback tensors (FLOAT in quant_meta, no .linear./.div. wrapper) are
    # skipped here; the base_state provides them unchanged.
    remapped: dict[str, torch.Tensor] = {}
    remapped_meta: dict[str, str] = {}
    mxfp4_layer_prefixes: set[str] = set()
    skipped: list[str] = []

    for key, tensor in quant_state.items():
        qtype = quant_meta.get(key, "FLOAT")

        if not _is_mxfp4_tensor(key, qtype):
            continue  # BF16 fallback — base_state already covers this tensor

        renamed = _apply_rename_dict(key)
        final_key = _strip_mxfp4_wrapper(renamed)

        if final_key.endswith(".quant_type"):
            skipped.append(key)
            continue

        remapped[final_key] = tensor
        remapped_meta[final_key] = qtype

        # Track MXFP4 layer prefixes (via their .weight keys) for ignored_layers.
        if final_key.endswith(".weight") and qtype.startswith("W4A4_MXFP4_DUALSCALE"):
            mxfp4_layer_prefixes.add(final_key[: -len(".weight")])

    if skipped:
        print(f"  Skipped {len(skipped)} metadata keys (quant_type markers): {skipped[:5]}")
    print(f"  {len(remapped)} MXFP4 tensors remapped, {len(mxfp4_layer_prefixes)} quantized layers")

    # Merge: base_state provides BF16 scaffold; MXFP4 tensors override their BF16 counterparts
    # and add the new scale tensors (weight_scale, weight_dual_scale, mul_scale).
    merged = {**base_state, **remapped}

    # Determine ignored_layers: BF16 fallback layer prefixes in vllm-omni parameter naming.
    ignored_layers = _collect_ignored_layers(merged, mxfp4_layer_prefixes)
    print(f"  {len(ignored_layers)} BF16 fallback layers → ignored_layers in config.json")

    # Save weights.
    out_weights = output_dir / "diffusion_pytorch_model.safetensors"
    save_file(merged, str(out_weights))
    print(f"  Saved {len(merged)} tensors → {out_weights}")

    # Save remapped quant metadata (for inspection / debugging).
    out_meta_path = output_dir / "quant_model_description.json"
    with open(out_meta_path, "w") as f:
        json.dump(remapped_meta, f, indent=2)

    # Inject quantization_config into config.json.
    src_config = original_transformer_dir / "config.json"
    if src_config.is_file():
        with open(src_config) as f:
            config = json.load(f)
        config["quantization_config"] = _build_quant_config(ignored_layers)
        out_config = output_dir / "config.json"
        with open(out_config, "w") as f:
            json.dump(config, f, indent=2)
        print(f"  Injected quantization_config (mxfp4_dualscale, {len(ignored_layers)} ignored_layers) → {out_config}")
    else:
        print(f"  WARNING: No config.json at {src_config}; quantization_config not injected.")


def _build_quant_config(ignored_layers: list[str]) -> dict[str, Any]:
    return {
        "quant_method": "mxfp4_dualscale",
        "is_checkpoint_serialized": True,
        "ignored_layers": ignored_layers,
    }


# ---------------------------------------------------------------------------
# Model-type helpers
# ---------------------------------------------------------------------------


def _get_transformer_dirs(model_type: str) -> list[str]:
    return ["transformer", "transformer_2"] if model_type in CASCADE_MODEL_TYPES else ["transformer"]


def _get_quant_subdir(model_type: str, quant_path: pathlib.Path, transformer_dir: str) -> pathlib.Path:
    if model_type in CASCADE_MODEL_TYPES:
        sub = "high_noise_model" if transformer_dir == "transformer" else "low_noise_model"
        return quant_path / sub
    return quant_path


# ---------------------------------------------------------------------------
# Main repack
# ---------------------------------------------------------------------------


def repack(
    model_type: str,
    original_model_path: pathlib.Path,
    quant_path: pathlib.Path,
    output_path: pathlib.Path,
) -> None:
    transformer_dirs = _get_transformer_dirs(model_type)

    print(f"Copying original model to {output_path} (skipping {transformer_dirs}) …")
    shutil.copytree(
        str(original_model_path),
        str(output_path),
        ignore=shutil.ignore_patterns(*transformer_dirs),
    )

    for tdir in transformer_dirs:
        q_subdir = _get_quant_subdir(model_type, quant_path, tdir)
        out_tdir = output_path / tdir
        orig_tdir = original_model_path / tdir
        print(f"\nConverting {tdir} (quant source: {q_subdir.name}) …")
        _convert_transformer(q_subdir, out_tdir, orig_tdir)

    print(f"\nDone. Merged model → {output_path}")
    print("\nRun inference (quantization auto-detected from config.json):")
    print(f"  python text_to_video.py --model {output_path}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--model-type", required=True, choices=SUPPORTED_MODEL_TYPES, help="Model variant.")
    parser.add_argument("--original-model", required=True, help="Original HF Diffusers model directory (BF16).")
    parser.add_argument("--quant-path", required=True, help="msModelSlim quantized weights directory.")
    parser.add_argument("--output-path", required=True, help="Output directory for merged model.")
    args = parser.parse_args()

    repack(
        model_type=args.model_type,
        original_model_path=pathlib.Path(args.original_model),
        quant_path=pathlib.Path(args.quant_path),
        output_path=pathlib.Path(args.output_path),
    )


if __name__ == "__main__":
    main()