#!/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-/ (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-).") 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__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()