quantize_cosmos3_i2v4step_streaming.py

#!/usr/bin/env python
"""
Stream-quantize NVIDIA Cosmos3-Super-Image2Video-4Step's transformer to weight-only
FP8 or NVFP4, WITHOUT ever materializing the full ~128 GB BF16 model.

This is the same-architecture sibling of quantize_cosmos3_super_streaming.py (the
recipe validated against nvidia/Cosmos3-Super): identical Cosmos3OmniTransformer
class, identical dims (64 layers, hidden_size 5120), so the empty-on-meta ->
mtq.quantize -> mtq.compress -> stream-load recipe and the SPARE_SUBSTRINGS list
carry over unchanged. Three things do NOT carry over -- this checkpoint is a DMD2
4-step distillation of Cosmos3-Super-Image2Video, and differs as follows:

  1. SCHEDULER IS FIXED, NOT SWAPPABLE. The base-model script (and
     serve_cosmos3_diffusers.py) replace the shipped scheduler with
     UniPCMultistepScheduler(flow_shift=...). This checkpoint ships
     FlowMatchEulerDiscreteScheduler with a baked-in fixed_step_sampler_config
     (sde sampling, t_list=[1.0, 0.9375, 0.8333, 0.625] -- literally 4 steps).
     Do NOT override it -- use pipe.scheduler as loaded.
  2. num_inference_steps / guidance_scale are NOT meaningfully configurable per
     NVIDIA's model card (CFG is distilled out; the checkpoint was trained for an
     exact 4-step sde schedule). CONFIRMED BY RUNNING IT: Cosmos3OmniPipeline
     (diffusers 0.39.0.dev0) does not know this checkpoint's scheduler carries a
     fixed_step_sampler_config -- an unpatched call silently ran the pipeline's
     35-step/guidance=6.0 default instead. render_from_memory() below patches the
     scheduler (_force_fixed_step_schedule) to force the correct 4-step t_list and
     passes guidance_scale=1.0 explicitly. See that function's docstring.
  3. This checkpoint is Image2Video-specific (unlike the T2I/T2V/I2V-omni base
     model), so the smoke render requires a conditioning IMAGE -- there is no
     text-only still-image path to fall back on. sound_tokenizer is also null in
     this repo (audio branch inactive), so the "audio_proj" spare substring is
     believed inert here, but is kept for parity/safety.

USAGE
-----
    python quantize_cosmos3_i2v4step_streaming.py --format fp8
    python quantize_cosmos3_i2v4step_streaming.py --format fp8 --smoke --image out.png

Outputs go to ./cosmos3-i2v4step-<fmt>/ (override with --export-dir).
"""

import argparse
import os

# Reduce allocator fragmentation on the big card (cheap, always-on).
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
from accelerate import init_empty_weights, load_checkpoint_in_model
from accelerate.utils import get_max_memory, infer_auto_device_map
from accelerate.utils.dataclasses import CustomDtype
from huggingface_hub import snapshot_download

import modelopt.torch.quantization as mtq
from modelopt.torch.export import export_hf_checkpoint

# Cosmos3 classes require diffusers built from git main (already installed in the venv).
from diffusers import Cosmos3OmniTransformer

SRC_REPO = "nvidia/Cosmos3-Super-Image2Video-4Step"

# ---------------------------------------------------------------------------
# Layers to KEEP IN BF16 (never quantize). Matched as substrings of module names.
# Unchanged from the base-model recipe -- same transformer class + dims, so the
# same layers are quality-sensitive: embeddings, norms, the bundled Qwen3 reasoner
# head, time/modality adapters, and the in/out projections.
# ---------------------------------------------------------------------------
SPARE_SUBSTRINGS = [
    "time_embedder",
    "proj_in",
    "proj_out",
    "lm_head",
    "embed",        # token / position embeddings
    "norm",         # layernorms / rmsnorms
    "audio_proj",   # audio modality adapter (inert here: sound_tokenizer is null)
]


def _is_spare(module_name: str) -> bool:
    return any(s in module_name for s in SPARE_SUBSTRINGS)


def build_quant_cfg(fmt: str) -> dict:
    """Return a WEIGHT-ONLY quant config for the chosen format. Same as the
    base-model recipe: FP8 per-tensor E4M3, or NVFP4 4-bit block-scale."""
    if fmt == "fp8":
        return {
            "quant_cfg": {
                "*weight_quantizer": {"num_bits": (4, 3), "axis": None, "enable": True},
                "*input_quantizer": {"enable": False},
                "*output_quantizer": {"enable": False},
                "*softmax_quantizer": {"enable": False},
            },
            "algorithm": "max",
        }
    elif fmt == "nvfp4":
        import copy

        base = getattr(mtq, "W4A16_NVFP4_CFG", None) or mtq.NVFP4_DEFAULT_CFG
        return copy.deepcopy(base)
    else:
        raise ValueError(f"Unknown format: {fmt!r}")


def enforce_weight_only_and_spare(model) -> tuple[int, int]:
    """Disable all activation quantizers, and weight quantizers on SPARE layers.
    Config-form agnostic (works for both the FP8 dict and modelopt's NVFP4 preset)."""
    n_spare = 0
    n_act = 0
    for name, module in model.named_modules():
        if not (name.endswith("_quantizer") and hasattr(module, "disable")):
            continue
        if name.endswith("weight_quantizer"):
            parent = name.rsplit(".", 1)[0]
            if _is_spare(parent):
                module.disable()
                n_spare += 1
        else:
            module.disable()
            n_act += 1
    return n_spare, n_act


def compressed_device_map(model, gpu_mem_fraction: float = 0.85) -> dict:
    """Build a device_map sized for the COMPRESSED weights (same as base-model recipe)."""
    max_memory = {k: v * gpu_mem_fraction for k, v in get_max_memory().items()}

    no_split = set()
    for name, module in model.named_modules():
        if name.endswith((".layers.0", ".blocks.0", ".transformer_blocks.0")):
            no_split.add(module.__class__.__name__)

    special_dtypes = {}
    for name, module in model.named_modules():
        if (
            hasattr(module, "weight")
            and hasattr(module, "weight_quantizer")
            and getattr(module.weight_quantizer, "is_enabled", True)
            and not getattr(module.weight_quantizer, "fake_quant", True)
        ):
            nb = module.weight_quantizer.num_bits
            if isinstance(nb, tuple):
                nb = nb[0] + nb[1] + 1
            special_dtypes[name + ".weight"] = CustomDtype.FP8 if nb == 8 else CustomDtype.INT4

    return infer_auto_device_map(
        model,
        max_memory=max_memory,
        no_split_module_classes=list(no_split),
        special_dtypes=special_dtypes,
    )


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--format", choices=["fp8", "nvfp4"], required=True,
                    help="Weight-only quantization format to produce.")
    ap.add_argument("--export-dir", default=None,
                    help="Output dir (default: ./cosmos3-i2v4step-<format>).")
    ap.add_argument("--gpu-mem-fraction", type=float, default=0.85,
                    help="Fraction of each GPU's memory accelerate may use for placement.")
    ap.add_argument("--smoke", action="store_true",
                    help="Validate by rendering a short clip from the in-memory model "
                         "(runs before export; output: cosmos3_i2v4step_<fmt>_validate.mp4).")
    ap.add_argument("--image", default="out.png",
                    help="Conditioning image for the --smoke render (this checkpoint is "
                         "Image2Video-only -- there is no text-only still path).")
    args = ap.parse_args()

    export_dir = args.export_dir or f"./cosmos3-i2v4step-{args.format}"
    os.makedirs(export_dir, exist_ok=True)

    print(f"[1/6] Resolving local checkpoint for {SRC_REPO} (transformer only)...")
    local_root = snapshot_download(SRC_REPO, allow_patterns=["transformer/*"])
    transformer_dir = os.path.join(local_root, "transformer")
    print(f"      transformer dir: {transformer_dir}")

    print("[2/6] Building EMPTY transformer on meta device (params on meta, buffers real)...")
    config = Cosmos3OmniTransformer.load_config(transformer_dir)
    with init_empty_weights(include_buffers=False):
        model = Cosmos3OmniTransformer.from_config(config)

    print(f"[3/6] Inserting quantizers ({args.format}) on the meta model...")
    quant_cfg = build_quant_cfg(args.format)
    mtq.quantize(model, quant_cfg)  # no forward_loop: weight-only needs no calibration
    n_spare, n_act = enforce_weight_only_and_spare(model)
    print(f"      weight-only: disabled {n_act} activation quantizers; "
          f"kept {n_spare} projection layers in BF16 (plus embeddings/norms/head)")

    print("[4/6] Setting up compressed parameter shapes (mtq.compress)...")
    try:
        mtq.compress(model, config=mtq.CompressConfig(quant_gemm=False))
    except (AttributeError, TypeError):
        mtq.compress(model)

    print("[5/6] Streaming BF16 shards into compressed form (this is the long step)...")
    device_map = compressed_device_map(model, args.gpu_mem_fraction)
    load_checkpoint_in_model(
        model,
        checkpoint=transformer_dir,
        device_map=device_map,
        dtype=torch.bfloat16,
    )

    n_fixed = 0
    for _, module in model.named_modules():
        for bname, buf in list(module._buffers.items()):
            if buf is not None and getattr(buf, "is_meta", False):
                module._buffers[bname] = torch.zeros(buf.shape, dtype=buf.dtype, device="cuda")
                n_fixed += 1
        for pname, par in list(module._parameters.items()):
            if par is not None and getattr(par, "is_meta", False):
                module._parameters[pname] = torch.nn.Parameter(
                    torch.zeros(par.shape, dtype=par.dtype, device="cuda"), requires_grad=False
                )
                n_fixed += 1
    if n_fixed:
        print(f"      materialized {n_fixed} residual meta tensors (disabled-quantizer scratch)")

    n_bytes = sum(p.numel() * p.element_size() for p in model.parameters() if p.device.type != "meta")
    n_bytes += sum(b.numel() * b.element_size() for b in model.buffers() if b.device.type != "meta")
    print(f"      live footprint: {n_bytes / 1e9:.1f} GB")

    if args.smoke:
        render_from_memory(model, args.format, args.image)

    print(f"[6/6] Exporting unified HF checkpoint to {export_dir} ...")
    with torch.inference_mode():
        export_hf_checkpoint(model, export_dir=export_dir)
    print(f"DONE. Quantized {args.format.upper()} checkpoint written to {export_dir}")


def _force_fixed_step_schedule(scheduler) -> bool:
    """CONFIRMED (not speculative): Cosmos3OmniPipeline.__call__ always calls
    scheduler.set_timesteps(num_inference_steps, device=device) with no sigmas
    passthrough (pipeline_cosmos3_omni.py:1529 in diffusers 0.39.0.dev0), so this
    checkpoint's scheduler_config.json fixed_step_sampler_config (t_list) is silently
    dropped -- an unpatched call falls back to the pipeline's num_inference_steps=35
    default (verified empirically: a real run produced a 35-step progress bar despite
    this being a 4-step DMD2-distilled checkpoint). Running any schedule other than the
    exact trained t_list is off-distribution for what the model was distilled to
    denoise, so patch THIS scheduler instance's set_timesteps to always use it,
    regardless of whatever num_inference_steps the pipeline passes in.

    ALSO disables stochastic_sampling (forces the deterministic ODE branch) -- CONFIRMED
    by direct A/B render, not speculative. This checkpoint ships stochastic_sampling=True
    (SDE sampling). Cosmos3's image-conditioning anchors frame 0 by zeroing the model's
    predicted velocity there (_mask_velocity_predictions), which only means "leave this
    position unchanged" under the DETERMINISTIC step (prev_sample = sample + dt*0 =
    sample). The SDE branch instead computes
        x0 = sample - current_sigma * model_output   # = sample when velocity is 0
        prev_sample = (1 - next_sigma) * x0 + next_sigma * randn_tensor(...)
    which re-noises the sample by next_sigma regardless of velocity -- there is no
    zero-velocity no-op in the SDE formula. Across this checkpoint's 4 steps that
    compounds to ~99.6% fresh noise by the final step: the conditioning image is
    destroyed into colorful static while the (non-anchored, genuinely denoised) motion
    frames still look like a plausible video. Disabling stochastic_sampling makes zero
    velocity a true no-op again, which an A/B render (same seed/image/prompt) confirmed
    fixes it outright. This is presumably NVIDIA's own reference runtime doing a masked
    resample this diffusers pipeline port doesn't implement; revisit if that lands
    upstream."""
    cfg = getattr(scheduler.config, "fixed_step_sampler_config", None)
    t_list = cfg.get("t_list") if isinstance(cfg, dict) else None
    if not t_list:
        print("[warn] no fixed_step_sampler_config.t_list on this scheduler; leaving set_timesteps unpatched")
        return False
    _orig = scheduler.set_timesteps

    def _patched(num_inference_steps=None, device=None, sigmas=None, mu=None, timesteps=None):
        return _orig(sigmas=list(t_list), device=device)

    scheduler.set_timesteps = _patched
    if getattr(scheduler.config, "stochastic_sampling", False):
        scheduler.register_to_config(stochastic_sampling=False)
        print("[patch] disabled stochastic_sampling (SDE re-noising corrupts the "
              "image-conditioned frame -- see docstring)")
    print(f"[patch] forced fixed {len(t_list)}-step sde schedule: {t_list}")
    return True


def render_from_memory(model, fmt: str, image_path: str):
    """Validate the quantization by rendering an image->video clip from the IN-MEMORY
    model. Does NOT swap the scheduler object itself (keeps FlowMatchEulerDiscreteScheduler
    with its stochastic_sampling config), but DOES patch its set_timesteps via
    _force_fixed_step_schedule -- see that function for why this is required rather than
    optional. guidance_scale=1.0 is passed explicitly: CFG is distilled out of this
    checkpoint (do_classifier_free_guidance is guidance_scale != 1.0 in the pipeline
    source), and the 6.0 pipeline default would silently re-enable it."""
    import gc

    print(f"\n[validate] Rendering an image->video clip from the in-memory {fmt.upper()} model "
          f"(conditioning image: {image_path})...")
    try:
        from PIL import Image

        from diffusers import Cosmos3OmniPipeline

        if not os.path.isfile(image_path):
            raise FileNotFoundError(
                f"conditioning image {image_path!r} not found -- pass --image with a real "
                "file; this checkpoint has no text-only still path to fall back on."
            )
        cond_image = Image.open(image_path).convert("RGB")

        model.to("cuda")

        # Dtype consistency at render time (same nudge as the base-model script): our
        # model mixes FP8/NVFP4 (compressed) weights with BF16 spare weights, while
        # diffusers computes the timestep sinusoidal fresh in FP32 each step.
        for _module in model.modules():
            for _bn, _buf in list(_module._buffers.items()):
                if _buf is not None and _buf.dtype == torch.float32:
                    _module._buffers[_bn] = _buf.to(torch.bfloat16)

        def _cast_inputs_bf16(_m, args):
            return tuple(
                a.to(torch.bfloat16)
                if torch.is_tensor(a) and a.is_floating_point() and a.dtype != torch.bfloat16
                else a
                for a in args
            )

        n_hooks = 0
        for _name, _module in model.named_modules():
            if "time_embedder" in _name and hasattr(_module, "linear_1"):
                _module.register_forward_pre_hook(_cast_inputs_bf16)
                n_hooks += 1
        print(f"[validate] dtype-safety: cast fp32 buffers to bf16, hooked {n_hooks} time-embedder(s)")

        # Pass OUR quantized transformer in so the pipeline does NOT reload it from the hub.
        # Deliberately do NOT touch pipe.scheduler -- the shipped FlowMatchEulerDiscreteScheduler
        # carries the fixed 4-step sde schedule for this checkpoint.
        pipe = Cosmos3OmniPipeline.from_pretrained(
            SRC_REPO,
            transformer=model,
            torch_dtype=torch.bfloat16,
            enable_safety_checker=False,
        )
        for name, comp in pipe.components.items():
            if name != "transformer" and isinstance(comp, torch.nn.Module):
                comp.to("cuda")

        _force_fixed_step_schedule(pipe.scheduler)

        prompt = (
            "The camera holds static as the scene continues naturally, with subtle "
            "motion and realistic physics."
        )
        with torch.inference_mode():
            # num_inference_steps is a no-op once _force_fixed_step_schedule has patched
            # set_timesteps (it always substitutes the checkpoint's own t_list), but is
            # passed for clarity. guidance_scale=1.0 disables CFG -- see docstring above.
            # 13 frames (4n+1) keeps the smoke test quick; bump for a real quality check.
            result = pipe(
                prompt=prompt,
                negative_prompt="",
                image=cond_image,
                num_frames=13,
                height=256,
                width=256,
                num_inference_steps=4,
                guidance_scale=1.0,
                generator=torch.Generator(device="cuda").manual_seed(1234),
            )
        out_path = f"cosmos3_i2v4step_{fmt}_validate.mp4"
        from diffusers.utils import export_to_video
        export_to_video(result.video, out_path, fps=24)
        print(f"[validate] Wrote {out_path}. Eyeball it for coherence with the conditioning image.")

        del pipe, result
        gc.collect()
        torch.cuda.empty_cache()
    except Exception as e:
        import traceback
        print(f"[validate] Render failed ({type(e).__name__}: {e}).")
        print("[validate] This does NOT affect the quantized weights; export still proceeds below.")
        print("[validate] If this is a TypeError about an unexpected keyword argument, that's real "
              "signal about this checkpoint's actual pipe() call shape -- adjust the call above "
              "(e.g. drop/rename kwargs) rather than assuming the defaults here were correct.")
        traceback.print_exc()


if __name__ == "__main__":
    main()