modelopt/onnx/quantization/quant_utils.py (3)

169-171: Consider inlining this simple calculation.

This helper function is only used once (in _pad) and performs a trivial calculation. Consider inlining it directly where used to reduce unnecessary abstraction.

-def _next_block_size_multiple(x: float, block_size: int) -> float:
-    return math.ceil(x / block_size) * block_size
-

 def _pad(w: np.ndarray, block_size: int, quantize_axis: int = 0) -> np.ndarray:
     """Pads `w` to next largest multiple of block_size, on quantize_axis."""
     assert quantize_axis <= len(w.shape), (
         f"incorrect quantize-axis {quantize_axis}, w-shape={w.shape}"
     )
     
     if w.shape[quantize_axis] % block_size == 0:
         return w
     
-    pad_width = (
-        _next_block_size_multiple(w.shape[quantize_axis], block_size) - w.shape[quantize_axis]
-    )
+    pad_width = math.ceil(w.shape[quantize_axis] / block_size) * block_size - w.shape[quantize_axis]

---

190-201: Consider using numpy slicing syntax for all axes.

The function handles only 2D arrays but could be made more general and consistent using numpy's ellipsis syntax.

 def _depad(w: np.ndarray, orig_shape: tuple, quantize_axis: int = 0) -> np.ndarray:
     """Depad quantize_axis to original shape."""
     if w.shape == orig_shape:
         return w
-    ans = None
-    if quantize_axis == 0:
-        ans = w[0 : orig_shape[0], ...]
-    elif quantize_axis == 1:
-        ans = w[..., 0 : orig_shape[1]]
-    else:
-        raise ValueError("Incorrect Quantize-axis: it must be 0 or 1 for a 2D array")
-    return ans
+    
+    if quantize_axis not in [0, 1]:
+        raise ValueError("Incorrect Quantize-axis: it must be 0 or 1 for a 2D array")
+    
+    slices = [slice(None)] * len(w.shape)
+    slices[quantize_axis] = slice(0, orig_shape[quantize_axis])
+    return w[tuple(slices)]

---

323-338: Add docstring details about return values.

The function's docstring should specify what the three return values represent.

 def quant_tensor(
     w: np.ndarray,
     block_size: int,
     quantize_axis: int = 0,
     alpha: float = 1.0,
     use_zero_point: bool = False,
     precision_info: dict[str, int] | None = None,
     name: str | None = None,
 ):
-    """Quantize a tensor using alpha etc. and return the quantized tensor."""
+    """Quantize a tensor using alpha etc. and return the quantized tensor.
+    
+    Returns:
+        tuple: A tuple containing:
+            - wq: The quantized weight tensor (np.ndarray)
+            - scale: The scale factors used for quantization (np.ndarray)
+            - zp: The zero-point values (np.ndarray or None if not using zero-point)
+    """
     scale, zp = find_scales(
         w, block_size, quantize_axis, alpha, use_zero_point, precision_info, name
     )
     wq = rtn(w, scale, block_size, quantize_axis, zp, precision_info, name)
     return wq, scale, zp

examples/windows/onnx_ptq/genai_llm/quantize.py (1)

599-613: Consider validating the int8_layers pattern format.

The int8_layers parameter accepts a comma-separated string of layer patterns, but there's no validation of the format. Consider adding basic validation to catch user errors early.

 parser.add_argument(
     "--int8_layers",
     type=str,
     default="",
     help=(
         "Comma-separated list of layer patterns to quantize to INT8 instead of INT4."
         "Example: 'layers.0,layers.1,lm_head'"
     ),
 )
+
+def validate_int8_layers(layers_str: str) -> bool:
+    """Validate the format of int8_layers string."""
+    if not layers_str:
+        return True
+    # Basic validation: check for valid characters and structure
+    import re
+    pattern = r'^[a-zA-Z0-9_.,\-]+$'
+    return bool(re.match(pattern, layers_str))

Then add validation after parsing:

if args.int8_layers and not validate_int8_layers(args.int8_layers):
    parser.error("Invalid format for --int8_layers. Use comma-separated layer patterns.")

modelopt/onnx/quantization/qdq_utils.py (3)

310-322: Add input validation for the attributes parameter.

The function modifies the attributes dictionary in place, which could cause unexpected side effects if the same dictionary is reused.

 def update_attributes(attrib: dict[str, Any] | None = None, is_per_channel: bool = False):
     """Update attribute dictionary for quantization nodes.
     
     If per-channel quantization is enabled, sets the 'axis' attribute to 1 and removes
     the 'block_size' attribute if present.
+    
+    Args:
+        attrib: Attribute dictionary to update (will be modified in place if not None)
+        is_per_channel: Whether to enable per-channel quantization
+    
+    Returns:
+        The updated attribute dictionary
     """
     if is_per_channel:
         if attrib is not None:
             attrib["axis"] = 1
             if "block_size" in attrib:
                 attrib.pop("block_size")
     return attrib

---

354-369: Consider extracting tensor dtype selection logic to a helper function.

The tensor dtype selection logic is duplicated in both insert_dq_nodes and insert_qdq_nodes. Consider extracting it to reduce duplication and improve maintainability.

+def get_tensor_dtype(precision_info: dict[str, int] | None, name: str, has_zero_point: bool) -> int:
+    """Get the appropriate tensor dtype based on precision info and zero point presence.
+    
+    Args:
+        precision_info: Dictionary mapping tensor names to bit widths
+        name: Name of the tensor
+        has_zero_point: Whether the tensor has a zero point
+    
+    Returns:
+        ONNX tensor data type constant
+    """
+    if precision_info and name in precision_info:
+        bit_width = precision_info[name]
+        if has_zero_point:
+            dtype_str = onnx_bit_dtype_unsigned_map[bit_width]
+        else:
+            dtype_str = onnx_bit_dtype_signed_map[bit_width]
+        return onnx_dtype_map[dtype_str]
+    else:
+        # Default to 4-bit
+        return onnx.TensorProto.INT4 if not has_zero_point else onnx.TensorProto.UINT4

 def _insert_helper(
     name: str,
     wq: np.ndarray,
     scale: np.ndarray,
     dq_nodes: dict[str, gs.Node],
     zp: np.ndarray,
     attrs: dict[str, Any] | None = None,
     precision_info: dict[str, int] | None = None,
     is_per_channel: bool = False,
 ):
     attrib = dict(attrs) if attrs is not None else None
-    if precision_info and name in precision_info:
-        tensor_dtype = (
-            onnx_dtype_map[onnx_bit_dtype_signed_map[precision_info[name]]]
-            if zp is None
-            else onnx_dtype_map[onnx_bit_dtype_unsigned_map[precision_info[name]]]
-        )
-        # do per-channel quantization for int8 as no support for int8 block-wise dq node
-        if precision_info[name] == 8:
-            # reshape scale to be per-channel
-            scale = scale.reshape(-1)
-            attrib = update_attributes(attrib, True)
-    else:
-        tensor_dtype = onnx.TensorProto.INT4 if zp is None else onnx.TensorProto.UINT4
+    tensor_dtype = get_tensor_dtype(precision_info, name, zp is not None)
+    
+    # do per-channel quantization for int8 as no support for int8 block-wise dq node
+    if precision_info and name in precision_info and precision_info[name] == 8:
+        # reshape scale to be per-channel
+        scale = scale.reshape(-1)
+        attrib = update_attributes(attrib, True)

---

361-364: Validate per‐channel scale size before flattening. Although scale.reshape(-1) preserves all elements, add an explicit check (e.g. assert scale.size == wq.shape[1]) before calling reshape to guard against mismatches when quantizing along axis 1.

modelopt/onnx/quantization/int4.py (1)

620-624: Use logger.warning instead of deprecated logger.warn.

-        logger.warn("Augmented ONNX model or external data file was not found")
+        logger.warning("Augmented ONNX model or external data file was not found")

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/onnx/quantization/quant_utils.py (1)

292-326: Numerical safety in rtn.

Add a guard to avoid tiny s after broadcast.

-    if zp is None:
+    # Safety: avoid division by tiny scales
+    if np.any(np.abs(s) < CLIP_MIN):
+        s = np.where(np.abs(s) < CLIP_MIN, CLIP_MIN, s)
+    if zp is None:
         maxq = 2 ** (num_bits - 1) - 1

modelopt/onnx/quantization/graph_utils.py (2)

629-667: Type hint misuse and minor robustness.

• cast("int", weight_tensor.data_type) uses a string literal as a type param; use cast(int, ...) or drop cast.
• Consider guarding for Gemm transA (currently TODO) or explicitly documenting unsupported cases.

-        gemm_io_type = cast("int", weight_tensor.data_type)
+        gemm_io_type = cast(int, weight_tensor.data_type)

---

669-696: Token-exact matcher looks good. Add fast exit and docstring type.

Add -> bool and early return for long patterns.

-def should_quantize_to_int8(layer_name: str, int8_layers: list[str]):
+def should_quantize_to_int8(layer_name: str, int8_layers: list[str]) -> bool:
@@
-    for pat in int8_layers:
+    for pat in int8_layers:
         needle = tokens(pat)
         n, m = len(hay), len(needle)
+        if m == 0 or m > n:
+            continue

modelopt/onnx/quantization/quant_utils.py (2)

160-177: update_block_size: clarify contract for w=None.

Function dereferences w unconditionally when block_size==-1 or num_bits==8. Either assert or handle None.

-    if block_size is not None and (block_size == -1 or num_bits == 8):
-        return w.shape[quantize_axis]
+    if block_size is not None and (block_size == -1 or num_bits == 8):
+        assert w is not None, "update_block_size requires `w` when per-channel/INT8 is requested"
+        return w.shape[quantize_axis]

---

196-198: Type hint for _next_block_size_multiple.

Return type is int in practice; adjust for readability/tools.

-def _next_block_size_multiple(x: float, block_size: int) -> float:
+def _next_block_size_multiple(x: float, block_size: int) -> int:

modelopt/onnx/quantization/qdq_utils.py (1)

50-63: ONNX 4-bit enums: confirm availability.

onnx.TensorProto.INT4/UINT4 are relatively new; some environments may lack them. Add a fallback or guard.

Option: probe availability once and raise a clear error.

-onnx_dtype_map = {
+onnx_dtype_map = {
@@
-    "INT4": onnx.TensorProto.INT4,
-    "UINT4": onnx.TensorProto.UINT4,
+    "INT4": getattr(onnx.TensorProto, "INT4", None),
+    "UINT4": getattr(onnx.TensorProto, "UINT4", None),
 }
+
+# Validate availability early
+if onnx_dtype_map["INT4"] is None or onnx_dtype_map["UINT4"] is None:
+    logger.error("This ONNX build does not support 4-bit dtypes (INT4/UINT4).")
+    # Optionally: raise, or fall back to INT8 path depending on product decision

Would you like me to wire a clean fallback to INT8 when 4-bit enums are missing?

modelopt/onnx/quantization/int4.py (2)

221-227: Pattern-vs-name expansion mismatch risk.

nodes_to_exclude is already expanded to concrete names before calling get_precision_info, which expects patterns (it internally expands again). This is benign but redundant and can mis-handle special regex chars in names.

Two options:

• Pass the original patterns into get_precision_info, or
• Change get_precision_info to accept already-expanded names and skip expansion.

---

1369-1376: Doc/type mismatch for int8_layers.

Docstring says “comma-separated list” but default shows [] and the code expects a string or None. Clarify.

-                - **int8_layers** (str): comma-separated list of layer patterns to quantize to INT8 instead of INT4.
-                                              Default: [].
+                - **int8_layers** (str | None): comma-separated list of layer patterns to quantize to INT8 instead of INT4.
+                                              Default: None.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

examples/windows/onnx_ptq/genai_llm/quantize.py (2)

364-366: Harden the gate that flips mixed-quant on when --layers_8bit is provided.

Trim whitespace and only flip when non-empty; also normalize the string once.

-    if args.layers_8bit:
-        args.enable_mixed_quant = True
+    if args.layers_8bit and args.layers_8bit.strip():
+        args.layers_8bit = args.layers_8bit.strip()
+        args.enable_mixed_quant = True

---

603-617: Tighten CLI help and document implied behavior.

Clarify that providing --layers_8bit implies mixed quant, and fix example punctuation.

-    parser.add_argument(
-        "--enable_mixed_quant",
+    parser.add_argument(
+        "--enable_mixed_quant",
         default=False,
         action="store_true",
         help=(
-            "Use default mixed quantization strategy: first 1/8, last 1/8, and every 3rd attn, "
-            "mlp layers quantized to 8 bits; others to 4 bits."
+            "Enable mixed quantization. Default strategy: first 1/8, last 1/8, and every 3rd "
+            "attn/mlp layer quantized to INT8; others to INT4."
         ),
     )
-    parser.add_argument(
+    parser.add_argument(
         "--layers_8bit",
         type=str,
         default="",
-        help=("Overrides default mixed quant strategy. Example: 'layers.0,lm_head'"),
+        help=("Overrides default strategy. Comma-separated patterns; providing this implies "
+              "--enable_mixed_quant. Example: 'layers.0,lm_head'."),
     )

modelopt/onnx/quantization/int4.py (4)

180-186: Avoid mutable default argument (nodes_to_exclude).

Using [] as a default risks cross-call mutation.

-def quantize_rtn(
+def quantize_rtn(
     onnx_model: onnx.ModelProto,
     block_size: int,
     dq_only: bool = False,
-    nodes_to_exclude: list[str] = [],
+    nodes_to_exclude: list[str] | None = None,
     **kwargs: Any,
 ) -> onnx.ModelProto:
@@
-    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)
+    nodes_to_exclude = nodes_to_exclude or []
+    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)

---

439-448: Same mutable default issue in _quantize_awq_clip.

-def _quantize_awq_clip(
+def _quantize_awq_clip(
     onnx_model: onnx.ModelProto,
     data_reader: CalibrationDataReader,
     use_external_data_format: bool,
     calibration_eps: list[str],
     block_size: int,
     force_fp16: bool = False,
-    nodes_to_exclude: list[str] = [],
+    nodes_to_exclude: list[str] | None = None,
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
     **kwargs: Any,
 ) -> onnx.ModelProto:
@@
-    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)
+    nodes_to_exclude = nodes_to_exclude or []
+    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)

---

959-972: Same mutable default issue in _quantize_awq_lite.

-def _quantize_awq_lite(
+def _quantize_awq_lite(
     onnx_model: onnx.ModelProto,
     data_reader: CalibrationDataReader,
     use_external_data_format: bool,
     calibration_eps: list[str],
     block_size: int,
     force_fp16: bool = False,
     enable_fast_path_using_high_sysram: bool = False,
     enable_weight_clipping: bool = False,
     use_zero_point: bool = False,
-    nodes_to_exclude: list[str] = [],
+    nodes_to_exclude: list[str] | None = None,
     input_shapes_profile: Sequence[dict[str, str]] | None = None,
     **kwargs: Any,
 ) -> onnx.ModelProto:
@@
-    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)
+    nodes_to_exclude = nodes_to_exclude or []
+    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)

---

604-604: Use logger.warning instead of deprecated logger.warn.

-        logger.warn("Augmented ONNX model or external data file was not found")
+        logger.warning("Augmented ONNX model or external data file was not found")

modelopt/onnx/quantization/qdq_utils.py (1)

311-325: Fail fast for unsupported bit-widths in get_tensor_dtype.

Avoid KeyError by validating num_bits and mapping lookups.

 def get_tensor_dtype(num_bits: int = 4, has_zero_point: bool = False) -> int:
@@
-    if has_zero_point:
-        dtype_str = onnx_bit_dtype_unsigned_map[num_bits]
-    else:
-        dtype_str = onnx_bit_dtype_signed_map[num_bits]
-    return onnx_dtype_map[dtype_str]
+    valid_bits = set(onnx_bit_dtype_unsigned_map) & set(onnx_bit_dtype_signed_map)
+    if num_bits not in valid_bits:
+        raise ValueError(f"Unsupported num_bits={num_bits}; expected one of {sorted(valid_bits)}")
+    dtype_str = (
+        onnx_bit_dtype_unsigned_map[num_bits]
+        if has_zero_point
+        else onnx_bit_dtype_signed_map[num_bits]
+    )
+    if dtype_str not in onnx_dtype_map:
+        raise ValueError(f"Unsupported dtype '{dtype_str}' in onnx_dtype_map")
+    return onnx_dtype_map[dtype_str]

modelopt/onnx/quantization/graph_utils.py (4)

707-714: Avoid mutable default (nodes_to_exclude) and normalize early.

-def get_layer_precision_mapping(
+def get_layer_precision_mapping(
     onnx_model: onnx.ModelProto,
     precision_pattern_8bit: str | None = None,
-    nodes_to_exclude: list[str] | None = [r"/lm_head"],
+    nodes_to_exclude: list[str] | None = None,
 ):
@@
-    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)
+    nodes_to_exclude = nodes_to_exclude or [r"/lm_head"]
+    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)

---

812-816: Same mutable default in get_precision_info.

-def get_precision_info(
+def get_precision_info(
     onnx_model: onnx.ModelProto,
-    nodes_to_exclude: list[str] | None = [r"/lm_head"],
+    nodes_to_exclude: list[str] | None = None,
     **kwargs: Any,
 ):
@@
-    precision_info = None
+    nodes_to_exclude = nodes_to_exclude or [r"/lm_head"]
+    precision_info = None

---

698-705: Slightly stricter validation using fullmatch.

Prevents accidental partial matches.

-    return bool(re.match(pattern, layers_str))
+    return bool(re.fullmatch(pattern, layers_str))

---

844-856: Consider re.fullmatch for pattern expansion (optional).

If callers pass anchored patterns it's fine; fullmatch would make intent explicit.

-        matched_node_names.extend([node.name for node in node_list if re.match(pattern, node.name)])
+        matched_node_names.extend([node.name for node in node_list if re.fullmatch(pattern, node.name)])

modelopt/onnx/quantization/quant_utils.py (1)

160-177: Guard update_block_size when w is required.

When block_size == -1 or num_bits == 8, w must be provided.

 def update_block_size(
     num_bits: int, block_size: int, quantize_axis: int = 0, w: np.ndarray = None
 ) -> int:
@@
-    if block_size is not None and (block_size == -1 or num_bits == 8):
-        return w.shape[quantize_axis]
+    if block_size is not None and (block_size == -1 or num_bits == 8):
+        assert w is not None, "Weight tensor 'w' is required to resolve per-channel block size"
+        return w.shape[quantize_axis]
     return block_size

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

modelopt/onnx/quantization/int4.py (3)

149-156: Keep zero‑point dtype consistent (don’t cast zp to weight dtype).

Pass zp unchanged to rtn. Casting zp to w.dtype can lead to subtle rounding differences; zp is an integer offset by definition.

Apply:

-                    qw = rtn(
-                        np.asarray(w),
-                        s,
-                        block_size_updated,
-                        quantize_axis=gather_quantize_axis,
-                        zp=zp if zp is None else zp.astype(w.dtype),
-                        num_bits=num_bits,
-                    )
+                    qw = rtn(
+                        np.asarray(w),
+                        s,
+                        block_size_updated,
+                        quantize_axis=gather_quantize_axis,
+                        zp=zp,
+                        num_bits=num_bits,
+                    )

---

1327-1333: Avoid mutable default for nodes_to_exclude in public API.

Use None default to prevent cross‑call mutation.

Apply:

-    nodes_to_exclude: list[str] | None = [r"/lm_head"],
+    nodes_to_exclude: list[str] | None = None,
@@
-    nodes_to_exclude = nodes_to_exclude or []
+    nodes_to_exclude = nodes_to_exclude or [r"/lm_head"]

---

97-101: Remove unused CLIP_MIN constant here (dup in quant_utils).

It’s duplicated and unused in this module.

Apply:

-# following min-value for clip is taken from AutoAWQ where zero-point based quantization is
-# supported and working
-CLIP_MIN = 1e-5

modelopt/onnx/quantization/graph_utils.py (2)

698-705: Regex is fine; consider re.fullmatch for clarity (optional).

Anchors already enforce full‑string match; switching to re.fullmatch would be clearer but not required.

---

707-711: Avoid mutable default list for nodes_to_exclude.

Prevent shared state across calls.

Apply:

-    nodes_to_exclude: list[str] | None = [r"/lm_head"],
+    nodes_to_exclude: list[str] | None = None,
@@
-    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude)
+    nodes_to_exclude = expand_node_names_from_patterns(graph, nodes_to_exclude or [r"/lm_head"])

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro