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

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

Based on the SGLang convert_wan_to_diffusers approach, extended to:
  - Inject `quantization_config` into each transformer's config.json so that
    vllm-omni's TransformerConfig.from_dict() auto-detects offline MXFP8 mode
    and selects NPUMxfp8LinearMethod (instead of the online/BF16 path).
  - Support MoE models (Wan2.2-T2V-A14B / I2V-A14B) with two transformers:
      quant_path/high_noise_model -> output/transformer
      quant_path/low_noise_model  -> output/transformer_2

Without the quantization_config injection, FP8 weights are passed to the online
quantization path (npu_dynamic_mx_quant), producing garbage output (visual artifacts).

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

Usage:
  python vllm_omni/quantization/tools/merge_mxfp8_checkpoint.py \\
      --model-type      Wan2.2-T2V-A14B \\
      --original-model  /path/to/Wan2.2-T2V-A14B-Diffusers \\
      --quant-path      /path/to/quant-output \\
      --output-path     /path/to/merged-Wan2.2-T2V-A14B-MXFP8
"""

from __future__ import annotations

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

from safetensors.torch import load_file, save_file

# ---------------------------------------------------------------------------
# Name mapping: quantization tool naming → diffusers naming
# ---------------------------------------------------------------------------

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 (original: 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", "Wan2.2-TI2V-5B"]

# MoE cascade models have high_noise_model + low_noise_model
CASCADE_MODEL_TYPES = {"Wan2.2-T2V-A14B", "Wan2.2-I2V-A14B"}

# quantization_config injected into each transformer's config.json.
# is_checkpoint_mxfp8_serialized=True triggers NPUMxfp8LinearMethod (offline path)
# instead of NPUMxfp8OnlineLinearMethod (online/BF16 path).
MXFP8_QUANT_CONFIG: dict[str, Any] = {
    "quant_method": "mxfp8",
    "is_checkpoint_mxfp8_serialized": True,
}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _get_transformer_dirs(model_type: str) -> list[str]:
    if model_type in CASCADE_MODEL_TYPES:
        return ["transformer", "transformer_2"]
    return ["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


def _remap_keys(state_dict: dict, quant_meta: dict) -> tuple[dict, dict]:
    """Apply TRANSFORMER_KEYS_RENAME_DICT to both state_dict and quant_meta."""
    new_state: dict = {}
    new_meta: dict = {}

    for key, tensor in state_dict.items():
        new_key = key
        for src, dst in TRANSFORMER_KEYS_RENAME_DICT.items():
            new_key = new_key.replace(src, dst)
        new_state[new_key] = tensor
        if key in quant_meta:
            new_meta[new_key] = quant_meta[key]

    return new_state, new_meta


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


def _load_quant_safetensors(directory: pathlib.Path) -> dict:
    try:
        return _load_safetensors(directory, "quant_model_weight*.safetensors")
    except FileNotFoundError:
        return _load_safetensors(directory, "*.safetensors")


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


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


def _convert_transformer(
    model_type: str,
    quant_subdir: pathlib.Path,
    output_dir: pathlib.Path,
    original_transformer_dir: pathlib.Path,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load original BF16 weights as the merge base so that tensors msModelSlim
    # omits (biases, norms, embeddings) are preserved without gaps.
    print(f"  Loading original BF16 weights from {original_transformer_dir} …")
    base_state_dict = _load_safetensors(original_transformer_dir)
    print(f"  {len(base_state_dict)} BF16 tensors loaded")

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

    # Rename keys from msModelSlim convention to diffusers convention
    quant_state_dict, quant_meta = _remap_keys(quant_state_dict, quant_meta)

    # Overlay: start from BF16 base, then apply all msModelSlim tensors on top.
    # Quantized weight tensors + their scale tensors override the BF16 originals.
    # Non-quantized tensors that msModelSlim emits (FLOAT in quant_meta) keep the
    # msModelSlim copy; any tensor msModelSlim omits is retained from the BF16 base.
    merged = {**base_state_dict, **quant_state_dict}

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

    # Save renamed quant metadata (optional, for reference)
    out_meta = output_dir / "quant_model_description.json"
    with open(out_meta, "w") as f:
        json.dump(quant_meta, f, indent=2)

    # Copy config.json from original transformer and inject quantization_config.
    # This is the critical step: TransformerConfig.from_dict() reads
    # config["quantization_config"]["quant_method"] to select the quant method,
    # and "is_checkpoint_mxfp8_serialized": true to select NPUMxfp8LinearMethod
    # (offline path) instead of NPUMxfp8OnlineLinearMethod (online/BF16 path).
    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"] = MXFP8_QUANT_CONFIG
        with open(output_dir / "config.json", "w") as f:
            json.dump(config, f, indent=2)
        print(f"  Injected quantization_config into config.json → {output_dir / 'config.json'}")
    else:
        print(f"  WARNING: No config.json found at {src_config}")


# ---------------------------------------------------------------------------
# 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)

    # Step 1: Copy original model, skipping transformer dirs (will be replaced).
    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),
    )

    # Step 2+: Convert each transformer.
    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(model_type, q_subdir, out_tdir, orig_tdir)

    print(f"\nDone. Merged model → {output_path}")
    print("\nRun inference (no --quantization flag needed; 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)
    parser.add_argument("--original-model", required=True, help="Original HF Diffusers model directory")
    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()