Move checkpoint_dir helpers from library to examples/llm_ptq

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

examples/llm_ptq/example_utils.py

@@ -15,6 +15,7 @@
 
 import copy
 import glob
+import hashlib
 import inspect
 import json
 import logging
@@ -854,3 +855,35 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod
         print(f"Successfully copied {len(copied_files)} custom model files to {export_path}")
     else:
         print("No custom model files found to copy")
+
+
+def needs_checkpoint_path_update(quant_cfg: dict) -> bool:
+    """Check if quant_cfg has a checkpoint_dir that should be auto-resolved to a unique subpath."""
+    algorithm = quant_cfg.get("algorithm")
+    if algorithm is None or isinstance(algorithm, str):
+        return False
+    return algorithm.get("checkpoint_dir") is not None
+
+
+def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> dict:
+    """Append a unique ``<model_name>_<config_hash>`` subdirectory to checkpoint_dir.
+
+    Allows a single recipe to be reused across models without checkpoint collisions.
+    Must only be called when :func:`needs_checkpoint_path_update` returns True.
+    """
+    algorithm = quant_cfg["algorithm"]
+    base_dir = algorithm["checkpoint_dir"]
+
+    name = model_path.rstrip("/")
+    if "/" in name and not os.path.isabs(name):
+        name = name.replace("/", "--")
+    else:
+        name = Path(name).name
+
+    config_hash = hashlib.sha256(
+        json.dumps(quant_cfg, sort_keys=True, default=str).encode()
+    ).hexdigest()[:8]
+
+    quant_cfg = copy.deepcopy(quant_cfg)
+    quant_cfg["algorithm"]["checkpoint_dir"] = os.path.join(base_dir, f"{name}_{config_hash}")
+    return quant_cfg

examples/llm_ptq/hf_ptq.py

@@ -34,6 +34,8 @@
     is_enc_dec,
     is_nemotron_vl,
     load_mtp_weights,
+    needs_checkpoint_path_update,
+    resolve_checkpoint_dir,
     run_nemotron_vl_preview,
 )
 from torch.utils.data import DataLoader
@@ -1105,11 +1107,6 @@ def quantize_main(
             quant_cfg = copy.deepcopy(quant_cfg)
             _set_kv_cache_constant_amax(quant_cfg["quant_cfg"])
 
-        from modelopt.torch.quantization.utils.layerwise_calib import (
-            needs_checkpoint_path_update,
-            resolve_checkpoint_dir,
-        )
-
         if needs_checkpoint_path_update(quant_cfg):
             quant_cfg = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path)
             print(f"Auto-resolved checkpoint_dir: {quant_cfg['algorithm']['checkpoint_dir']}")

modelopt/torch/quantization/utils/layerwise_calib.py

@@ -506,42 +506,6 @@ def _save_layer(
     _write_manifest(checkpoint_dir, idx, num_layers)
 
 
-def needs_checkpoint_path_update(quant_cfg: dict) -> bool:
-    """Check if quant_cfg has a checkpoint_dir that should be auto-resolved to a unique subpath."""
-    algorithm = quant_cfg.get("algorithm")
-    if algorithm is None or isinstance(algorithm, str):
-        return False
-    return algorithm.get("checkpoint_dir") is not None
-
-
-def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> dict:
-    """Append a unique ``<model_name>_<config_hash>`` subdirectory to checkpoint_dir.
-
-    Allows a single recipe to be reused across models without checkpoint collisions.
-    Must only be called when :func:`needs_checkpoint_path_update` returns True.
-    """
-    import copy
-    import hashlib
-    from pathlib import Path
-
-    algorithm = quant_cfg["algorithm"]
-    base_dir = algorithm["checkpoint_dir"]
-
-    name = model_path.rstrip("/")
-    if "/" in name and not os.path.isabs(name):
-        name = name.replace("/", "--")
-    else:
-        name = Path(name).name
-
-    config_hash = hashlib.sha256(
-        json.dumps(quant_cfg, sort_keys=True, default=str).encode()
-    ).hexdigest()[:8]
-
-    quant_cfg = copy.deepcopy(quant_cfg)
-    quant_cfg["algorithm"]["checkpoint_dir"] = os.path.join(base_dir, f"{name}_{config_hash}")
-    return quant_cfg
-
-
 def detect_resume_point(checkpoint_dir: str) -> tuple[int, dict] | None:
     """Detect where to resume from an existing checkpoint directory.