Merge branch 'main' into jingyux/diffusion.export-fixed

Author: jingyu-ml

CHANGELOG.rst

@@ -11,6 +11,7 @@ NVIDIA Model Optimizer Changelog (Linux)
 **New Features**
 
 - Add standalone type inference option (``--use_standalone_type_inference``) in ONNX AutoCast as an alternative to ONNX's ``infer_shapes``. This experimental feature performs type-only inference without shape inference, useful as a workaround when shape inference fails or to avoid unnecessary shape inference overhead.
+- Add support for Kimi K2 Thinking model quantization from the original int4 checkpoint.
 
 0.41 (2026-01-19)
 ^^^^^^^^^^^^^^^^^

docs/source/reference/2_security.rst

@@ -0,0 +1,78 @@
+Security Considerations
+=======================
+
+Overview
+--------
+
+NVIDIA Model Optimizer (ModelOpt) is a library used to optimize ML models and
+may load and process user-provided artifacts (models, weights, configs,
+calibration data) and their dependencies. Secure deployment depends on how you
+source artifacts, validate inputs, and harden the environment where ModelOpt
+runs.
+
+What to Be Aware Of
+-------------------
+
+**Untrusted model and data inputs**
+
+- Models, weights, configs and data may be malicious or corrupted.
+
+**Deserialization and code-execution risks**
+
+- Unsafe deserialization can lead to arbitrary code execution if fed untrusted
+  inputs.
+- Avoid using serialization formats/settings that can deserialize arbitrary
+  objects.
+
+**Input validation and resource exhaustion**
+
+- Large or malformed inputs can trigger crashes or excessive CPU/GPU/memory use.
+- Missing size/type checks can increase DoS risk.
+
+**Data in transit and at rest**
+
+- If fetching models or dependencies over the network, insecure transport can
+  enable tampering.
+- Stored artifacts, logs, and caches may contain sensitive data.
+
+**Logging and observability**
+
+- Logs may inadvertently contain sensitive inputs, paths, tokens, or proprietary
+  model details.
+- Overly verbose logs can leak operational and security-relevant information.
+
+**Supply chain and third-party components**
+
+- Dependencies may include known vulnerabilities or be compromised.
+- Third-party plugins/components loaded at runtime may not have the same
+  security assurances.
+
+Example Security Approaches
+---------------------------
+
+**Artifact integrity**
+
+- Only load artifacts from trusted sources.
+- Prefer signed artifacts; verify signatures before loading.
+
+**Safe parsing and deserialization**
+
+- Prefer safer storage formats (avoid object deserialization for untrusted
+  inputs).
+- Avoid ``pickle``, ``torch.load()`` with untrusted weights, or YAML
+  ``unsafe_load``.
+- Treat any unverified artifact as untrusted and block/guard its loading.
+
+**Hardening and least privilege**
+
+- Run with least privilege and isolate workloads.
+
+**Data protection**
+
+- Encrypt sensitive data at rest; use TLS 1.3 for data in transit.
+- Never hardcode or log credentials.
+
+**Resilience**
+
+- Validate inputs and enforce limits (file size, timeouts, quotas,..).
+- Keep OS, containers, and dependencies patched; scan for known vulnerabilities.

examples/llm_ptq/README.md

@@ -109,6 +109,7 @@ Please reference our [framework scripts](#framework-scripts) and our [docs](http
 | QWen3 MOE, Next <sup>6</sup> | ✅ | - | - | - | ✅ |
 | QwQ | ✅ | - | - | - | ✅ |
 | DeepSeek V3, R1, V3.1, V3.2<sup>7</sup> | - | - | - | - | ✅ |
+| Kimi K2 | - | - | - | - | ✅ |
 | T5 | ✅ | ✅ | ✅ | ✅ | - |
 | Whisper | ✅ | ❌ | ❌ | ❌ | - |
 

examples/llm_ptq/example_utils.py

@@ -187,6 +187,11 @@ def build_quant_cfg(
         quant_cfg["quant_cfg"]["model*.*attn*k_proj*"] = {"enable": False}
         quant_cfg["quant_cfg"]["model*.*attn*v_proj*"] = {"enable": False}
 
+    if model_type == "deepseek":
+        # Disable MLA quantization for accuracy.
+        quant_cfg["quant_cfg"]["*self_attn.q*"] = {"enable": False}
+        quant_cfg["quant_cfg"]["*self_attn.kv*"] = {"enable": False}
+
     return quant_cfg
 
 
@@ -346,6 +351,17 @@ def get_model(
                 device_map=device_map,
                 **model_kwargs,
             )
+        elif (
+            hasattr(hf_config, "quantization_config")
+            and hf_config.quantization_config.get("format", None) == "pack-quantized"
+        ):
+            torch_dtype = getattr(hf_config, "torch_dtype", torch.bfloat16)
+            model = AutoModelForCausalLM.from_pretrained(
+                ckpt_path,
+                device_map="auto",
+                trust_remote_code=trust_remote_code,
+                torch_dtype=torch_dtype,
+            )
         else:
             architecture = hf_config.architectures[0]
 
@@ -366,9 +382,9 @@ def get_model(
                 from_config = auto_model_module._from_config
 
             with init_empty_weights():
-                # When computing the device_map, assuming half precision by default,
+                # When computing the device_map, assuming bfloat16 precision by default,
                 # unless specified by the hf_config.
-                torch_dtype = getattr(hf_config, "torch_dtype", torch.float16)
+                torch_dtype = getattr(hf_config, "torch_dtype", torch.bfloat16)
                 model_kwargs2 = model_kwargs.copy()
                 if auto_model_module != AutoModelForCausalLM:
                     model_kwargs2.pop("trust_remote_code", None)

examples/llm_ptq/hf_ptq.py

@@ -575,7 +575,10 @@ def pre_quantize(
     ][0:1]
 
     # Generate preview before quantization
-    if is_nemotron_vl_model and tokenizer is not None:
+    if model_type == "deepseek":
+        # DeepSeek generation may go OOM, so we skip it
+        generated_ids_before_ptq = None
+    elif is_nemotron_vl_model and tokenizer is not None:
         generated_ids_before_ptq = run_nemotron_vl_preview(
             full_model,
             tokenizer,
@@ -618,7 +621,9 @@ def post_quantize(
     # Run some samples
     torch.cuda.empty_cache()
     generated_ids_after_ptq = None
-    if model_type != "llama4" and not is_nemotron_vl_model:
+    if generated_ids_before_ptq is None:
+        pass
+    elif model_type != "llama4" and not is_nemotron_vl_model:
         # Our fake quantizer may not be fully compatible with torch.compile.
         generated_ids_after_ptq = full_model.generate(preview_input_ids, max_new_tokens=100)
     elif is_nemotron_vl_model and tokenizer is not None:

examples/llm_ptq/requirements.txt

@@ -1,3 +1,4 @@
+compressed-tensors==0.12.0
 fire
 flash-attn>=2.6.0
 rouge_score>=0.1.2

examples/llm_ptq/scripts/huggingface_example.sh

@@ -53,9 +53,9 @@ esac
 IFS=","
 for qformat in $QFORMAT; do
     case $qformat in
-    fp8 | fp8_pc_pt | fp8_pb_wo | int8_wo | int8_sq | int4_awq | w4a8_awq | fp16 | bf16 | nvfp4 | nvfp4_awq | w4a8_nvfp4_fp8 | w4a8_mxfp4_fp8) ;;
+    fp8 | fp8_pc_pt | fp8_pb_wo | int8_wo | int8_sq | int4_awq | w4a8_awq | fp16 | bf16 | nvfp4 | nvfp4_awq | w4a8_nvfp4_fp8 | w4a8_mxfp4_fp8 | nvfp4_mlp_only) ;;
     *)
-        echo "Unknown quant argument: Expected one of: [fp8, fp8_pc_pt, fp8_pb_wo, int8_wo, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8]" >&2
+        echo "Unknown quant argument: Expected one of: [fp8, fp8_pc_pt, fp8_pb_wo, int8_wo, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, nvfp4_mlp_only]" >&2
         exit 1
         ;;
     esac

modelopt/torch/export/layer_utils.py

@@ -345,7 +345,8 @@ def is_moe(module: nn.Module) -> bool:
 
 def is_quantlinear(module: nn.Module) -> bool:
     """Returns whether the module is a quantized linear layer."""
-    return "QuantLinear" in type(module).__name__ and "lora" not in type(module).__name__.lower()
+    name = type(module).__name__
+    return ("QuantLinear" in name or "QuantCompressedLinear" in name) and "lora" not in name.lower()
 
 
 def dup_kv_weight(v: torch.Tensor, head_size: int, num_head: int, tp_size: int) -> torch.Tensor:

modelopt/torch/export/quant_utils.py

@@ -881,7 +881,13 @@ def postprocess_state_dict(
         "v_bmm_quantizer._bias_value": "v_proj.v_bias",
         "input_quantizer._pre_quant_scale": "pre_quant_scale",
     }
-    skip_keys = ["output_quantizer", "_amax", "_bias_value", "input_quantizer._pre_quant_scale"]
+    skip_keys = [
+        "output_quantizer",
+        "_amax",
+        "_bias_value",
+        "input_quantizer._pre_quant_scale",
+        "weight_shape",
+    ]
 
     # For modelopt-trained LoRA models, we need to remove the base_layer prefix from the keys for deployment
     if is_modelopt_qlora:

modelopt/torch/export/unified_export_hf.py

@@ -515,6 +515,8 @@ def _export_quantized_weight(
     if weight_scale is not None:
         sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale)
 
+    torch.cuda.empty_cache()
+
 
 def _process_quantized_modules(
     model: nn.Module,
@@ -550,6 +552,8 @@ def _process_quantized_modules(
         if is_modelopt_qlora and (hasattr(sub_module, "base_layer")):
             continue
 
+        if hasattr(sub_module, "weight_packed"):
+            sub_module.unpack_weight()
         if get_quantization_format(sub_module) != QUANTIZATION_NONE:
             if is_quantlinear(sub_module):
                 with fsdp2_aware_weight_update(model, sub_module, reshard=False):
@@ -948,6 +952,10 @@ def export_hf_checkpoint(
 
             hf_quant_config = convert_hf_quant_config_format(hf_quant_config)
 
+        # Remove hf_quantizer from model so post_state_dict can be exported.
+        if getattr(model, "hf_quantizer", None) is not None:
+            model.hf_quantizer = None
+
         # Save model
         model.save_pretrained(
             export_dir, state_dict=post_state_dict, save_modelopt_state=save_modelopt_state