Rename layerwise config fields and enable layerwise on experts-only recipe

- use_layerwise -> layerwise, checkpoint_dir -> layerwise_checkpoint_dir
- Enable layerwise calibration + checkpointing on nvfp4_experts_only-fp8_kv recipe
- Add layerwise_checkpoint_dir to nvfp4_default-none_kv_gptq recipe

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

@@ -858,21 +858,21 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod
 
 
 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."""
+    """Check if quant_cfg has a layerwise_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
+    return algorithm.get("layerwise_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.
+    """Append a unique ``<model_name>_<config_hash>`` subdirectory to layerwise_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"]
+    base_dir = algorithm["layerwise_checkpoint_dir"]
 
     name = model_path.rstrip("/")
     if "/" in name and not os.path.isabs(name):
@@ -885,5 +885,7 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> dict:
     ).hexdigest()[:8]
 
     quant_cfg = copy.deepcopy(quant_cfg)
-    quant_cfg["algorithm"]["checkpoint_dir"] = os.path.join(base_dir, f"{name}_{config_hash}")
+    quant_cfg["algorithm"]["layerwise_checkpoint_dir"] = os.path.join(
+        base_dir, f"{name}_{config_hash}"
+    )
     return quant_cfg

examples/llm_ptq/hf_ptq.py

@@ -1109,7 +1109,9 @@ def quantize_main(
 
         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']}")
+            print(
+                f"Auto-resolved layerwise_checkpoint_dir: {quant_cfg['algorithm']['layerwise_checkpoint_dir']}"
+            )
 
         if args.qformat in QUANT_CFG_CHOICES:
             mono_quantize(

modelopt/torch/quantization/config.py

@@ -1217,7 +1217,7 @@ class QuantizeAlgorithmConfig(ModeloptBaseConfig):
         ),
     )
 
-    use_layerwise: bool = ModeloptField(
+    layerwise: bool = ModeloptField(
         default=False,
         title="Enable layerwise (layer-by-layer) calibration.",
         description=(
@@ -1227,11 +1227,11 @@ class QuantizeAlgorithmConfig(ModeloptBaseConfig):
         ),
     )
 
-    checkpoint_dir: str | None = ModeloptField(
+    layerwise_checkpoint_dir: str | None = ModeloptField(
         default=None,
         title="Checkpoint directory for layerwise calibration.",
         description=(
-            "If set together with use_layerwise=True, per-layer checkpoints are saved to this "
+            "If set together with layerwise=True, per-layer checkpoints are saved to this "
             "directory during calibration. On restart, calibration resumes from the last "
             "completed layer."
         ),

modelopt/torch/quantization/mode.py

@@ -222,8 +222,8 @@ def wrapped_calib_func(
     """
     kwargs = config.model_dump()
     method = kwargs.pop("method")
-    layerwise = kwargs.pop("use_layerwise", False)
-    checkpoint_dir = kwargs.pop("checkpoint_dir", None)
+    layerwise = kwargs.pop("layerwise", False)
+    checkpoint_dir = kwargs.pop("layerwise_checkpoint_dir", None)
     if method is not None and "awq" in method:
         # For backward compatibility
         kwargs["algorithm"] = method

modelopt/torch/quantization/model_calib.py

@@ -1646,12 +1646,12 @@ def gptq(
 ):
     """GPTQ quantization.
 
-    Works in two modes depending on ``use_layerwise`` in the config:
+    Works in two modes depending on ``layerwise`` in the config:
 
-    * **Layerwise** (``use_layerwise=True``): ``layerwise_calibrate`` calls this
+    * **Layerwise** (``layerwise=True``): ``layerwise_calibrate`` calls this
       function once per decoder layer with updated activations, producing more
       accurate Hessian estimates.
-    * **Non-layerwise** (``use_layerwise=False``): called once on the full model.
+    * **Non-layerwise** (``layerwise=False``): called once on the full model.
       All layers are quantized in parallel from the original activations.
 
     Per-module steps:

modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yml

@@ -15,7 +15,7 @@
 
 metadata:
   recipe_type: ptq
-  description: NVFP4 MLP/MoE weight only (W4A16), FP8 KV cache, max calibration.
+  description: NVFP4 W4A4, FP8 KV cache, max calibration.
 quantize:
   algorithm: max
   quant_cfg:

modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yml

@@ -15,11 +15,12 @@
 
 metadata:
   recipe_type: ptq
-  description: NVFP4 weight and activation (W4A4), gptq sequential calibration.
+  description: NVFP4 weight and activation (W4A4), gptq layerwise calibration.
 quantize:
   algorithm:
     method: gptq
-    use_sequential: true
+    layerwise: true
+    layerwise_checkpoint_dir: output/layerwise_ckpts/
   quant_cfg:
     - quantizer_name: '*'
       enable: false

modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yml

@@ -15,9 +15,12 @@
 
 metadata:
   recipe_type: ptq
-  description: NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, max calibration.
+  description: NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, max layerwise calibration.
 quantize:
-  algorithm: max
+  algorithm:
+    method: max
+    layerwise: true
+    layerwise_checkpoint_dir: output/layerwise_ckpts/
   quant_cfg:
     - quantizer_name: '*'
       enable: false

tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py

@@ -100,20 +100,20 @@ def _make_cpu_offloaded_model(tmp_path, num_hidden_layers=3):
 
 
 def _make_layerwise_cfg(base_cfg):
-    """Add use_layerwise=True to a quant config's algorithm field."""
+    """Add layerwise=True to a quant config's algorithm field."""
     cfg = copy.deepcopy(base_cfg)
     algo = cfg.get("algorithm", "max")
     if isinstance(algo, str):
-        cfg["algorithm"] = {"method": algo, "use_layerwise": True}
+        cfg["algorithm"] = {"method": algo, "layerwise": True}
     else:
-        algo["use_layerwise"] = True
+        algo["layerwise"] = True
     return cfg
 
 
 def _make_layerwise_checkpoint_cfg(base_cfg, checkpoint_dir):
-    """Add use_layerwise=True and checkpoint_dir to a quant config's algorithm field."""
+    """Add layerwise=True and layerwise_checkpoint_dir to a quant config's algorithm field."""
     cfg = _make_layerwise_cfg(base_cfg)
-    cfg["algorithm"]["checkpoint_dir"] = checkpoint_dir
+    cfg["algorithm"]["layerwise_checkpoint_dir"] = checkpoint_dir
     return cfg
 
 
@@ -274,14 +274,14 @@ def _make_multi_offload_model():
 def _make_gptq_sequential_cfg(base_cfg):
     """Create a sequential GPTQ config from a base quantization config."""
     cfg = copy.deepcopy(base_cfg)
-    cfg["algorithm"] = {"method": "gptq", "use_layerwise": True}
+    cfg["algorithm"] = {"method": "gptq", "layerwise": True}
     return cfg
 
 
 def _make_gptq_sequential_checkpoint_cfg(base_cfg, checkpoint_dir):
     """Create a sequential GPTQ config with checkpoint dir."""
     cfg = _make_gptq_sequential_cfg(base_cfg)
-    cfg["algorithm"]["checkpoint_dir"] = checkpoint_dir
+    cfg["algorithm"]["layerwise_checkpoint_dir"] = checkpoint_dir
     return cfg
 
 

tests/gpu/torch/quantization/test_fsdp2.py

@@ -184,7 +184,7 @@ def _test_layerwise_calibrate_fsdp2(rank, size):
         # Reference: non-FSDP layerwise calibration
         ref_model = copy.deepcopy(model)
         seq_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG)
-        seq_cfg["algorithm"] = {"method": "max", "use_layerwise": True}
+        seq_cfg["algorithm"] = {"method": "max", "layerwise": True}
         mtq.quantize(ref_model, seq_cfg, lambda m: m(inputs))
         output_ref = ref_model(inputs)