fix(hf_ptq): scope layerwise export and resume to the one exporting pass

A sequential algorithm list made export ownership ambiguous. Detection read the
first layerwise block, so an exporting block behind a non-exporting one missed
the opt-in and exported to the recipe's placeholder; deriving the resume dir had
the mirror bug, letting another pass's explicit checkpoint_dir stand in for the
exporting one's, and giving every block one shared directory when none was set.

Export finalizes shards as calibration walks the layers, so only the final pass
can own them: layerwise_export_block() enforces exactly one export_dir on the
last entry and every export-driven helper reads through it. resolve_checkpoint_dir
now resolves each block against its own base, so two passes cannot collide on one
manifest.

Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>

Author: Fridah-nv

examples/hf_ptq/example_utils.py

@@ -1152,20 +1152,52 @@ def _layerwise_checkpoint_dir(algorithm) -> str | None:
     )
 
 
+def layerwise_export_block(algorithm) -> dict | None:
+    """The one ``layerwise`` block that owns per-layer export, or None.
+
+    Export finalizes each layer's shard during calibration, so a later pass would change
+    the model after its checkpoint was written: exactly one entry may set ``export_dir``,
+    and it must be the last.
+    """
+    entries = algorithm if isinstance(algorithm, list) else [algorithm]
+    exporting = [
+        (i, e["layerwise"])
+        for i, e in enumerate(entries)
+        if isinstance(e, dict)
+        and isinstance(e.get("layerwise"), dict)
+        and e["layerwise"].get("export_dir") is not None
+    ]
+    if not exporting:
+        return None
+    if len(exporting) > 1:
+        raise ValueError(
+            f"{len(exporting)} algorithm entries set layerwise.export_dir; only one "
+            "calibration pass can own the exported checkpoint."
+        )
+    index, block = exporting[0]
+    if index != len(entries) - 1:
+        raise ValueError(
+            f"layerwise.export_dir is set on algorithm entry {index} of {len(entries)}; it "
+            "must be the last, since a later pass would change the model after its shards "
+            "were written."
+        )
+    return block
+
+
 def default_layerwise_resume_dir(quant_cfg: dict, export_path: str) -> tuple[dict, bool]:
     """Derive ``layerwise.checkpoint_dir`` from ``export_path`` when unset.
 
     A sibling, not a child: nothing deletes the resume state, so inside ``export_path`` it
     would ship in the checkpoint. An explicit path is left alone.
     """
-    if _layerwise_checkpoint_dir(quant_cfg.get("algorithm")) is not None:
-        return quant_cfg, False
-
     quant_cfg = copy.deepcopy(quant_cfg)
-    blocks = _layerwise_blocks(quant_cfg.get("algorithm"))
-    for block in blocks:
-        block["checkpoint_dir"] = export_path.rstrip("/") + ".layerwise_resume"
-    return quant_cfg, bool(blocks)
+    # The exporting block specifically: another pass's explicit checkpoint_dir says nothing
+    # about where this one resumes from.
+    block = layerwise_export_block(quant_cfg.get("algorithm"))
+    if block is None or block.get("checkpoint_dir") is not None:
+        return quant_cfg, False
+    block["checkpoint_dir"] = export_path.rstrip("/") + ".layerwise_resume"
+    return quant_cfg, True
 
 
 def needs_checkpoint_path_update(quant_cfg: dict) -> bool:
@@ -1182,8 +1214,7 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str]
     Returns ``(updated_quant_cfg, resolved_path)`` so the caller can log or
     reference the resolved path without re-deriving the dict shape.
     """
-    base_dir = _layerwise_checkpoint_dir(quant_cfg["algorithm"])
-    assert base_dir is not None  # guaranteed by needs_checkpoint_path_update
+    assert needs_checkpoint_path_update(quant_cfg), "no layerwise.checkpoint_dir to resolve"
 
     name = model_path.rstrip("/")
     if "/" in name and not os.path.isabs(name):
@@ -1192,11 +1223,19 @@ def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> tuple[dict, str]
         name = Path(name).name
 
     config_hash = hashlib.sha256(json.dumps(quant_cfg, default=str).encode()).hexdigest()[:8]
-    resolved = os.path.join(base_dir, f"{name}_{config_hash}")
+    suffix = f"{name}_{config_hash}"
 
     quant_cfg = copy.deepcopy(quant_cfg)
+    # Each pass keeps its own base, so two layerwise passes cannot resolve onto one manifest.
+    exporting = layerwise_export_block(quant_cfg.get("algorithm"))
+    resolved = None
     for block in _layerwise_blocks(quant_cfg.get("algorithm")):
-        block["checkpoint_dir"] = resolved
+        if block.get("checkpoint_dir") is None:
+            continue
+        block["checkpoint_dir"] = os.path.join(block["checkpoint_dir"], suffix)
+        if resolved is None or block is exporting:
+            resolved = block["checkpoint_dir"]
+    assert resolved is not None  # needs_checkpoint_path_update found one above
     return quant_cfg, resolved
 
 
@@ -1210,22 +1249,14 @@ def set_layerwise_export_dir(quant_cfg: dict, export_path: str) -> dict:
     """
     quant_cfg = copy.deepcopy(quant_cfg)
     algorithm = quant_cfg.get("algorithm")
-    # Detection accepts one algorithm or a list, so the retarget must too.
-    retargeted = 0
-    for entry in algorithm if isinstance(algorithm, list) else [algorithm]:
-        # Only entries that already opted in: writing export_dir into a layerwise entry
-        # that did not ask for it would switch per-layer export on behind the user's back.
-        layerwise = entry.get("layerwise") if isinstance(entry, dict) else None
-        if isinstance(layerwise, dict) and layerwise.get("export_dir") is not None:
-            layerwise["export_dir"] = export_path
-            retargeted += 1
-
-    if not retargeted:
+    block = layerwise_export_block(algorithm)
+    if block is None:
         raise ValueError(
             "layerwise export is enabled but no layerwise.export_dir was found to retarget "
             f"in algorithm={algorithm!r}. The exported shards would go to the recipe's "
             "placeholder path instead of --export_path."
         )
+    block["export_dir"] = export_path
     return quant_cfg
 
 

examples/hf_ptq/hf_ptq.py

@@ -1200,33 +1200,35 @@ def quantize_main(
         aq_config = None
         fixed_quantize_config = None
 
-    def _layerwise_cfg(obj):
-        """The recipe's ``layerwise`` block, or None.
+    def _layerwise_cfgs(obj) -> list:
+        """The recipe's ``layerwise`` blocks, in algorithm order.
 
         An algorithm parsed from YAML arrives as a plain dict, while the deprecated
         ``--auto_quantize_*`` path builds config objects, so both shapes reach here.
         """
         if isinstance(obj, ModelOptPTQRecipe):
-            return _layerwise_cfg(obj.quantize.algorithm)
+            return _layerwise_cfgs(obj.quantize.algorithm)
         if isinstance(obj, ModelOptAutoQuantizeRecipe):
-            return _layerwise_cfg(obj.quantize.algorithm) if obj.quantize is not None else None
+            return _layerwise_cfgs(obj.quantize.algorithm) if obj.quantize is not None else []
         if isinstance(obj, list):
-            return next((cfg for cfg in map(_layerwise_cfg, obj) if cfg is not None), None)
-        if isinstance(obj, dict):
-            return obj.get("layerwise")
-        return getattr(obj, "layerwise", None)
+            return [cfg for entry in obj for cfg in _layerwise_cfgs(entry)]
+        cfg = obj.get("layerwise") if isinstance(obj, dict) else getattr(obj, "layerwise", None)
+        return [cfg] if cfg is not None else []
 
     def _layerwise_get(cfg, key, default=None):
         if cfg is None:
             return default
         return cfg.get(key, default) if isinstance(cfg, dict) else getattr(cfg, key, default)
 
-    layerwise_cfg = _layerwise_cfg(recipe)
-    is_layerwise = bool(_layerwise_get(layerwise_cfg, "enable", False))
+    layerwise_cfgs = _layerwise_cfgs(recipe)
+    is_layerwise = any(_layerwise_get(cfg, "enable", False) for cfg in layerwise_cfgs)
 
     # Setting layerwise.export_dir is the switch; the value is replaced with --export_path
-    # below, the way resolve_checkpoint_dir already rewrites layerwise.checkpoint_dir.
-    args.layerwise_export = _layerwise_get(layerwise_cfg, "export_dir") is not None
+    # below, the way resolve_checkpoint_dir already rewrites layerwise.checkpoint_dir. Any
+    # entry may hold it -- set_layerwise_export_dir settles which one legally owns it.
+    args.layerwise_export = any(
+        _layerwise_get(cfg, "export_dir") is not None for cfg in layerwise_cfgs
+    )
     if args.layerwise_export:
         if isinstance(recipe, ModelOptAutoQuantizeRecipe):
             # Only the mono-quantize path retargets export_dir and runs the refusals;

tests/examples/hf_ptq/test_example_utils.py

@@ -560,3 +560,86 @@ def from_pretrained(*args, **kwargs):
     example_utils.get_model("checkpoint", device="cpu", trust_remote_code=trust_remote_code)
 
     assert used["path"] == ("bundled" if expect_bundled_code else "builtin")
+
+
+def _layerwise(**kwargs):
+    return {"enable": True, **kwargs}
+
+
+def _blocks(quant_cfg):
+    algorithm = quant_cfg["algorithm"]
+    entries = algorithm if isinstance(algorithm, list) else [algorithm]
+    return [e["layerwise"] for e in entries if "layerwise" in e]
+
+
+@pytest.mark.parametrize(
+    ("algorithm", "expected"),
+    [
+        pytest.param(
+            {"method": "max", "layerwise": _layerwise(export_dir="/placeholder")},
+            ["/out.layerwise_resume"],
+            id="single-entry",
+        ),
+        pytest.param(
+            [
+                {"method": "awq", "layerwise": _layerwise()},
+                {"method": "max", "layerwise": _layerwise(export_dir="/placeholder")},
+            ],
+            [None, "/out.layerwise_resume"],
+            id="only-the-exporting-entry",
+        ),
+        pytest.param(
+            [
+                {"method": "awq", "layerwise": _layerwise(checkpoint_dir="/theirs")},
+                {"method": "max", "layerwise": _layerwise(export_dir="/placeholder")},
+            ],
+            ["/theirs", "/out.layerwise_resume"],
+            id="another-entrys-explicit-path-is-not-this-ones",
+        ),
+    ],
+)
+def test_default_layerwise_resume_dir_targets_the_exporting_entry(algorithm, expected):
+    """Only the pass that exports gets a derived resume dir, and only if it lacks one."""
+    updated, changed = example_utils.default_layerwise_resume_dir({"algorithm": algorithm}, "/out")
+
+    assert [b.get("checkpoint_dir") for b in _blocks(updated)] == expected
+    assert changed is True
+
+
+def test_resolve_checkpoint_dir_keeps_each_entrys_base():
+    """Two layerwise passes must not resolve onto one manifest."""
+    algorithm = [
+        {"method": "awq", "layerwise": _layerwise(checkpoint_dir="/theirs")},
+        {"method": "max", "layerwise": _layerwise(checkpoint_dir="/ours", export_dir="/ph")},
+    ]
+
+    updated, resolved = example_utils.resolve_checkpoint_dir({"algorithm": algorithm}, "/m/Model")
+
+    theirs, ours = (b["checkpoint_dir"] for b in _blocks(updated))
+    assert theirs.startswith("/theirs/") and ours.startswith("/ours/")
+    assert theirs != ours
+    # The exporting pass owns the path the caller reports.
+    assert resolved == ours
+
+
+@pytest.mark.parametrize(
+    ("algorithm", "match"),
+    [
+        pytest.param(
+            [
+                {"layerwise": _layerwise(export_dir="/a")},
+                {"layerwise": _layerwise(export_dir="/b")},
+            ],
+            "only one calibration pass",
+            id="two-exporting-entries",
+        ),
+        pytest.param(
+            [{"layerwise": _layerwise(export_dir="/a")}, {"method": "max"}],
+            "must be the last",
+            id="a-later-pass-would-change-the-model",
+        ),
+    ],
+)
+def test_set_layerwise_export_dir_refuses_ambiguous_ownership(algorithm, match):
+    with pytest.raises(ValueError, match=match):
+        example_utils.set_layerwise_export_dir({"algorithm": algorithm}, "/out")