#!/usr/bin/env python """ Drop-in loader for the weight-only NVFP4 / FP8 (NVIDIA ModelOpt) Cosmos3-Super checkpoint saved in the ROUND-TRIPPABLE format (modelopt_state.pth present -- see repackage_for_hf.py). Current diffusers + accelerate + modelopt treat this path as experimental; each shim below works around a specific, source-verified version-skew gap. None modify the checkpoint: 1. enable_huggingface_checkpointing() Registers ModelOpt's HF handlers so `from_pretrained` restores the quantized module structure from modelopt_state.pth before weights load. 2. set_module_tensor_to_device patch (parameter materialization) The modelopt_state restore runs inside diffusers' meta-device init, so each quantized weight is a QTensorWrapper whose storage is a META tensor carrying dequant metadata. `param.data = real` can't cross meta<->real, and accelerate's rebuild both rejects requires_grad and discards metadata. We REPLACE the parameter with a fresh wrapper around the loaded bytes + the existing metadata. Two extra duties here: - payload dtype restore: diffusers casts floating params to torch_dtype during load when no hf_quantizer is present (model_loading_utils.py -- their own TODO flags float8). FP8 payloads are floating and arrive cast to bf16; we cast back to the wrapper's payload dtype (exact: every e4m3fn value round-trips through bf16). NVFP4's uint8 payload is never cast, so this is a no-op there. - direct-to-GPU materialization: staging payloads on CPU would put the whole packed model in system RAM (FP8: ~64 GB on a 32 GB box -> OOM-killed; NVFP4 only survived because uncast tensors stay mmap-backed). Wrappers go straight to `materialize_device`. 3. Post-restore quantizer re-disable modelopt_state replays the QUANT CONFIG, not imperative `.disable()` calls made after quantize. The NVFP4 build used NVFP4_DEFAULT_CFG (activation quantization ON in-config) with activations disabled imperatively -- so the restored model comes back with ~1806 dynamic fake-quant activation quantizers active: ~10x slower, fatter, and quantizing activations the validated regime never quantized. We re-apply weight-only + spare disabling after load. (FP8's config had the disables baked in; this is then a no-op.) 4. Full bf16 normalization (the validated serve regime) Cosmos3OmniTransformer pins time_embedder to fp32 via _keep_in_fp32_modules (transformer_cosmos3.py:297); the quantized model runs all-bf16. Cast floating buffers and non-wrapper floating params to bf16, retarget wrapper dequant dtype to bf16, and register bf16 input-cast pre-hooks on the time embedders. 5. NVIDIAModelOptQuantizer.create_quantized_param patch Only relevant if a checkpoint carries an embedded quantization_config (hf_quantizer path): requires_grad only for float/complex tensors. Kept for robustness. Usage (library): from load_cosmos3_modelopt import load_pipe pipe = load_pipe("YOUR_HF_USERNAME/Cosmos3-Super-nvfp4") # or a local dir # NOTE: a bare pipe(prompt) call renders the pipeline DEFAULT: a 189-frame 720x1280 # video (~8 s at 24 fps) -- not a still. For a single-image smoke test, be explicit: r = pipe("a red cube on a table", height=1024, width=1024, num_frames=1, num_inference_steps=50, guidance_scale=4.0) r.video[0].save("out.png") # .video is the list of PIL frames; [0] is the image Usage (interactive -- lands in a REPL with `pipe` loaded): CUDA_VISIBLE_DEVICES=0 python -i load_cosmos3_modelopt.py ./cosmos3-super-nvfp4-hf """ import os os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import torch SPARE_SUBSTRINGS = [ "time_embedder", "proj_in", "proj_out", "lm_head", "embed", "norm", "audio_proj", ] def _qtensor_wrapper_cls(): try: from modelopt.torch.quantization.qtensor.base_qtensor import QTensorWrapper return QTensorWrapper except Exception: return None def _patch_modelopt_quantizer() -> None: """diffusers pre-quantized hf_quantizer branch: only float/complex tensors may require grad.""" import diffusers.quantizers.modelopt.modelopt_quantizer as moq if getattr(moq.NVIDIAModelOptQuantizer, "_cqp_patched", False): return _orig = moq.NVIDIAModelOptQuantizer.create_quantized_param def _cqp(self, model, param_value, param_name, target_device, *args, **kwargs): if self.pre_quantized: module, tname = moq.get_module_from_name(model, param_name) needs_grad = param_value.is_floating_point() or param_value.is_complex() module._parameters[tname] = torch.nn.Parameter( param_value.to(device=target_device), requires_grad=needs_grad ) return return _orig(self, model, param_value, param_name, target_device, *args, **kwargs) moq.NVIDIAModelOptQuantizer.create_quantized_param = _cqp moq.NVIDIAModelOptQuantizer._cqp_patched = True def _patch_qtensor_loading() -> None: """Materialize restored (meta) QTensorWrapper params: replace the parameter object with a fresh wrapper around the loaded bytes, restoring the payload dtype (undoes diffusers' float cast on FP8) and landing directly on the target device (keeps payloads out of RAM).""" import diffusers.models.model_loading_utils as mlu if getattr(mlu, "_qtw_inplace_patched", False): return QTensorWrapper = _qtensor_wrapper_cls() if QTensorWrapper is None: return # different modelopt layout; nothing to patch _orig = mlu.set_module_tensor_to_device stats = {"materialized": 0, "target_device": None} def _patched(model, tensor_name, device, value=None, *args, **kwargs): if value is not None: module, leaf = model, tensor_name if "." in tensor_name: mod_path, leaf = tensor_name.rsplit(".", 1) try: module = model.get_submodule(mod_path) except AttributeError: module = None if module is not None: cur = getattr(module, "_parameters", {}).get(leaf) if isinstance(cur, QTensorWrapper): tgt = stats.get("target_device") or device module._parameters[leaf] = QTensorWrapper( value.to(device=tgt, dtype=cur.data.dtype), # exact cast-back for fp8 metadata=dict(cur.metadata), ) stats["materialized"] += 1 return None return _orig(model, tensor_name, device, value=value, *args, **kwargs) mlu.set_module_tensor_to_device = _patched mlu._qtw_inplace_patched = True mlu._qtw_stats = stats def _enforce_weight_only(transformer) -> None: """Re-apply the validated weight-only + spare regime: the state replay re-enables any quantizers that were disabled imperatively after quantize (NVFP4 default cfg has activation quantization ON in-config).""" n_act = n_spare = 0 for name, m in transformer.named_modules(): if not (name.endswith("_quantizer") and hasattr(m, "disable")): continue if name.endswith("weight_quantizer"): if any(s in name for s in SPARE_SUBSTRINGS): if getattr(m, "is_enabled", False): n_spare += 1 m.disable() else: if getattr(m, "is_enabled", False): n_act += 1 m.disable() print(f"[load] re-disabled quantizers the state replay re-enabled: " f"{n_act} activation, {n_spare} spare-weight") def _apply_dtype_nudges(transformer) -> None: """Reproduce the validated all-bf16 serve regime (all no-ops where already aligned).""" QTensorWrapper = _qtensor_wrapper_cls() or () n_buf = n_par = n_meta = 0 for m in transformer.modules(): for bn, buf in list(m._buffers.items()): if buf is not None and buf.is_floating_point() and buf.dtype != torch.bfloat16: m._buffers[bn] = buf.to(torch.bfloat16) n_buf += 1 for _, p in transformer.named_parameters(): if isinstance(p, QTensorWrapper): d = p.metadata.get("dtype") if isinstance(d, torch.dtype) and d.is_floating_point and d != torch.bfloat16: p.metadata["dtype"] = torch.bfloat16 # dequant target only; payload untouched n_meta += 1 continue # packed payloads: never cast if p.is_floating_point() and p.dtype != torch.bfloat16: p.data = p.data.to(torch.bfloat16) n_par += 1 print(f"[load] normalized to bf16: {n_par} params, {n_buf} buffers; " f"retargeted {n_meta} dequant dtypes") def _cast_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 ) for name, m in transformer.named_modules(): if "time_embedder" in name and hasattr(m, "linear_1"): m.register_forward_pre_hook(_cast_bf16) def load_pipe( model_id_or_path: str, *, torch_dtype=torch.bfloat16, enable_safety_checker: bool = False, device: str = "cuda", materialize_device: str | None = "cuda", # packed weights stream straight here (RAM stays low) **kwargs, ): """Load a ModelOpt-quantized Cosmos3-Super pipeline with all load-time fixes applied.""" from diffusers import Cosmos3OmniPipeline from modelopt.torch.opt import enable_huggingface_checkpointing enable_huggingface_checkpointing() # must run before from_pretrained _patch_modelopt_quantizer() _patch_qtensor_loading() import diffusers.models.model_loading_utils as mlu if hasattr(mlu, "_qtw_stats"): mlu._qtw_stats["materialized"] = 0 mlu._qtw_stats["target_device"] = materialize_device pipe = Cosmos3OmniPipeline.from_pretrained( model_id_or_path, torch_dtype=torch_dtype, enable_safety_checker=enable_safety_checker, **kwargs, ) n_mat = getattr(mlu, "_qtw_stats", {}).get("materialized", 0) print(f"[load] materialized {n_mat} packed quantized weight tensors") transformer = getattr(pipe, "transformer", None) if transformer is not None: _enforce_weight_only(transformer) _apply_dtype_nudges(transformer) pipe = pipe.to(device) QTensorWrapper = _qtensor_wrapper_cls() if QTensorWrapper is not None and transformer is not None: n_live = sum(1 for p in transformer.parameters() if isinstance(p, QTensorWrapper)) print(f"[load] {n_live} quantized weight wrappers active after move to {device}") if n_mat and not n_live: print("[load] WARNING: wrappers were lost during .to() -- do not render; report this") return pipe if __name__ == "__main__": import sys path = sys.argv[1] if len(sys.argv) > 1 else "./cosmos3-super-nvfp4-hf" print(f"[load] loading {path} ...") pipe = load_pipe(path) print("[load] OK -- `pipe` is ready.") print(" NOTE: bare pipe(prompt) renders a 189-frame 720x1280 video (pipeline default).") print(" single-still smoke test:") print(" r = pipe('a red cube on a table', height=1024, width=1024, num_frames=1,") print(" num_inference_steps=50, guidance_scale=4.0); r.video[0].save('out.png')")