#!/usr/bin/env python """ Stream-quantize NVIDIA Cosmos3-Super's transformer to weight-only FP8 or NVFP4, WITHOUT ever materializing the full ~128 GB BF16 model. WHY THIS EXISTS --------------- The diffusers <-> ModelOpt quantize-on-load path materializes the full BF16 model on the GPU before compressing, which needs ~128 GB of VRAM for a 64B model and won't spill to CPU. That can't fit a 96 GB card. This script instead replicates ModelOpt's own `init_quantized_weights` (the engine behind `--low_memory_mode` in their llm_ptq example), which is the tool NVIDIA built for exactly "I can RUN the FP8 model but can't QUANTIZE it naively": 1. Build the transformer EMPTY on the meta device (zero real memory). 2. mtq.quantize(...) -> insert quantizers (weight-only; no model execution). 3. mtq.compress(...) -> set up REAL compressed (FP8/NVFP4) parameter shapes. 4. infer a device_map from those *compressed* sizes. 5. load_checkpoint_and_dispatch(...) -> stream the BF16 shards straight into compressed form, one shard at a time. Per-tensor weight scales are computed from each weight as it lands. The full BF16 is NEVER resident. Peak memory ~= compressed size + one shard: FP8 ~65 GB -> fits the 96 GB RTX 6000 Pro alone, comfortably. NVFP4 ~36 GB -> fits with enormous margin. Both formats are WEIGHT-ONLY, so neither runs a calibration forward pass -- the 64B model never executes. The only difference between them is the quant config; ModelOpt does the FP8 per-tensor scaling and the NVFP4 4-bit block-scale packing internally, so you get a fair FP8-vs-NVFP4 comparison from one script. Output is the ModelOpt "unified HF checkpoint" (safetensors + hf_quant_config.json), loadable by vLLM / TensorRT-LLM / diffusers. USAGE ----- python quantize_cosmos3_super_streaming.py --format fp8 python quantize_cosmos3_super_streaming.py --format nvfp4 # add --smoke to attempt a tiny render from the exported checkpoint afterward. Outputs go to ./cosmos3-super-/ (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" # --------------------------------------------------------------------------- # Layers to KEEP IN BF16 (never quantize). Matched as substrings of module names. # Mirrors the Cosmos3-Nano-FP8 community recipe + what --inspect showed for Super: # keep embeddings, norms, the bundled Qwen3 reasoner head, time/modality adapters, # and the in/out projections. These are tiny relative to the 64 transformer blocks # but quality-sensitive, so they cost almost nothing to leave in BF16. # # NOTE: substrings are deliberately specific. "proj_in"/"proj_out" do NOT match the # attention "add_q_proj"/"to_out"/"down_proj" etc. (different names), so the 64 MMDiT # blocks' attention + MLP linears still get quantized. # --------------------------------------------------------------------------- SPARE_SUBSTRINGS = [ "time_embedder", "proj_in", "proj_out", "lm_head", "embed", # token / position embeddings "norm", # layernorms / rmsnorms "audio_proj", # audio modality adapter ] 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. We start from ModelOpt's weight-only presets so the heavy lifting (the FP8 per-tensor scale and the NVFP4 E2M1 + block-scale packing) is ModelOpt's code, not ours. We only ensure activations stay off (weight-only) and normalize to a dict so the spare-exclusion below is unambiguous. """ if fmt == "fp8": # Per-tensor E4M3 weights, everything else in BF16. This is the exact dict # that previously passed "Inserted 2709 quantizers" under mtq. 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": # NVFP4: E2M1 4-bit weights, block_size 16, FP8 (E4M3) block scales + a FP32 # per-tensor scale. We base this on whichever NVFP4 config your installed # modelopt actually ships: newer versions have a weight-only W4A16_NVFP4_CFG, # older ones only NVFP4_DEFAULT_CFG (weights + activations). # # IMPORTANT: the weight-only disables must be baked INTO THE CONFIG, not just # applied imperatively afterwards. modelopt_state (written by save_pretrained # for the drop-in HF repo) replays the CONFIG on restore -- imperative # .disable() calls made after quantize are not captured, so a checkpoint built # from a bare NVFP4_DEFAULT_CFG restores with ~1806 dynamic activation # quantizers active (~10x slower per step). enforce_weight_only_and_spare() # below still runs as belt-and-braces for the live model. import copy base = getattr(mtq, "W4A16_NVFP4_CFG", None) or mtq.NVFP4_DEFAULT_CFG cfg = copy.deepcopy(base) cfg.setdefault("quant_cfg", {}) cfg["quant_cfg"]["*input_quantizer"] = {"enable": False} cfg["quant_cfg"]["*output_quantizer"] = {"enable": False} cfg["quant_cfg"]["*softmax_quantizer"] = {"enable": False} for s in SPARE_SUBSTRINGS: cfg["quant_cfg"][f"*{s}*weight_quantizer"] = {"enable": False} return cfg else: raise ValueError(f"Unknown format: {fmt!r}") def enforce_weight_only_and_spare(model) -> tuple[int, int]: """Make the model strictly WEIGHT-ONLY and keep SPARE layers in BF16. Runs after mtq.quantize (on the meta model) and before mtq.compress. Walks every inserted quantizer -- they appear as leaf modules whose name ends in '_quantizer': - any NON-weight quantizer (input/output/softmax/bmm activation) -> DISABLED, so activations stay BF16 regardless of which base config we started from. - a weight quantizer on a SPARE layer (embeddings/norms/head/adapters) -> DISABLED, so those weights stay BF16. This is config-form agnostic (works whether quant_cfg was a dict or a list), which is why we can feed it either the FP8 dict or modelopt's NVFP4 preset unchanged. Returns (spare_weight_quantizers_disabled, activation_quantizers_disabled). """ 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. Adapted from ModelOpt's init_quantized_weights.get_model_device_map: tell accelerate that each compressed weight is FP8 (8-bit) or INT4 (4-bit) so the map reflects ~65 GB / ~36 GB, not the 128 GB BF16 footprint. The result will place essentially everything on GPU 0 (it fits), but GPU 1 + CPU stay available as spill if a future/larger model needs them. """ max_memory = {k: v * gpu_mem_fraction for k, v in get_max_memory().items()} # Treat the first transformer block's class as un-splittable so a single block # isn't torn across devices (keeps attention math on one device). 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): # e.g. (4,3) for FP8, (2,1) for NVFP4 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-super-).") 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 one 1024x1024 image from the in-memory model " "(runs before export; output: cosmos3_super__validate.png).") args = ap.parse_args() export_dir = args.export_dir or f"./cosmos3-super-{args.format}" os.makedirs(export_dir, exist_ok=True) print(f"[1/6] Resolving local checkpoint for {SRC_REPO} (transformer only)...") # Pull just the transformer subfolder's config + safetensors shards locally. 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) # include_buffers=False is load-bearing: parameters still go to meta (the ~128 GB # we're avoiding), but buffers (rotary frequencies, masks, etc.) get REAL storage. # If buffers were meta too, accelerate's final dispatch -> model.to(device) crashes # with "Cannot copy out of meta tensor". mtq.quantize runs *after* this context, so # the quantizer scale buffers are real as well -> nothing meta survives to dispatch. 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)...") # quant_gemm=False matches ModelOpt's low-memory loader; export handles the rest. try: mtq.compress(model, config=mtq.CompressConfig(quant_gemm=False)) except (AttributeError, TypeError): # Older modelopt: compress takes no CompressConfig. 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 weights into the compressed (meta) model WITHOUT accelerate's dispatch step. # The all-in-one load_checkpoint_and_dispatch finishes with model.to(device), which # crashes on the leftover meta _amax scratch buffers of the ~1,800 DISABLED quantizers # (never written by the load). We load here, materialize those residual meta tensors # ourselves, then export. The real compressed weights + scales are filled by this call. load_checkpoint_in_model( model, checkpoint=transformer_dir, device_map=device_map, dtype=torch.bfloat16, ) # Materialize any residual meta buffers/params as zeros on GPU so nothing meta reaches # export. These belong to disabled (unused) quantizers, so zeros are inert; the enabled # weight quantizers' scales were already computed during the load above. 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)") # Footprint report (compressed weights + buffers actually resident) 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") # Validate BEFORE export: render from the pristine post-load model. export_hf_checkpoint # may mutate quantizer state (e.g. the QKV amax fusion), so we eyeball the real weights # first. render_from_memory frees its own GPU memory before export reuses the card. if args.smoke: render_from_memory(model, args.format) 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 render_from_memory(model, fmt: str): """Validate the quantization by rendering one image from the IN-MEMORY model. This renders directly from the transformer we just quantized -- no reload. That's the correct torch validation path: diffusers' torch round-trip uses mto.save/mto.restore, whereas what we export is the *deployment* unified checkpoint (for vLLM-Omni / TRT-LLM). Because mtq.compress used quant_gemm=False, the QuantLinears dequantize on the fly, so the rendered image reflects exactly the FP8/NVFP4 weight rounding -- which is what we want to eyeball. Best-effort: any failure here is reported but does not block export. """ import gc print(f"\n[validate] Rendering a 1024x1024 image from the in-memory {fmt.upper()} model...") try: from diffusers import Cosmos3OmniPipeline from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler # Consolidate everything onto GPU 0: the compressed weights are already there; this # pulls the small from_config buffers (rotary, etc.) over too. Safe now that the # materialize-meta step left nothing on meta. model.to("cuda") # Dtype consistency at render time. Our model mixes FP8 (compressed) weights with BF16 # spare weights, while diffusers computes the timestep sinusoidal fresh in FP32 each # step. That FP32 tensor then meets the BF16 time-embedder weights -> # "mat1 and mat2 must have the same dtype". NVIDIA's uniform-BF16 run never sees this; # our mixed model needs two nudges (runtime-only -- the exported weights are untouched): # (a) cast any FP32 buffers (from the empty-init) to BF16, leaving FP8 weights alone; # (b) cast time-embedder inputs to BF16 at the FP32->BF16 boundary. 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; # it only fetches the small components (VAE, scheduler, tokenizer). The reasoner tower # lives inside the transformer, so nothing large is double-loaded -> fits the 96 GB card. pipe = Cosmos3OmniPipeline.from_pretrained( SRC_REPO, transformer=model, torch_dtype=torch.bfloat16, enable_safety_checker=False, # skip the guardrail model for a local check ) pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=3.0) for name, comp in pipe.components.items(): if name != "transformer" and isinstance(comp, torch.nn.Module): comp.to("cuda") prompt = ( "A medium shot of a modern robotics research laboratory with white walls and a gray " "floor. A robotic arm with a metallic finish is mounted on a clean white workbench, " "its gripper positioned above a row of small colored objects. A large monitor on the " "wall behind displays a software interface, brightly lit by overhead lights." ) with torch.inference_mode(): result = pipe( prompt=prompt, negative_prompt="", num_frames=1, # single frame -> still image height=1024, width=1024, num_inference_steps=50, guidance_scale=4.0, generator=torch.Generator(device="cuda").manual_seed(1234), ) out_path = f"cosmos3_super_{fmt}_validate.png" result.video[0].save(out_path) print(f"[validate] Wrote {out_path}. Eyeball it for coherence; compare fp8 vs nvfp4 " f"(same prompt + seed 1234, so differences are purely the format).") # Free the pipeline's extra components/activations before export reuses the GPU. 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.") traceback.print_exc() if __name__ == "__main__": main()