Allow searcher ckpt dir for per-rank ckpt files (#1091)

### What does this PR do?

Type of change: minor improvement <!-- Use one of the following: Bug
fix, new feature, new example, new tests, documentation. -->

Currently we modify user's `path.pth` to `path<rank>.pth` which may be
confusing. Instead we now add alternative to provide `path/` and store
`path/rank<rank>.pth` per-rank searcher state files

### Testing
<!-- Mention how have you tested your change if applicable. -->

### Before your PR is "*Ready for review*"

Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)
and your commits are signed (`git commit -s -S`).

Make sure you read and follow the [Security Best
Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors)
(e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(...,
weights_only=False)`, `pickle`, etc.).

- Is this change backward compatible?: ✅ <!--- If ❌, explain why. -->
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A <!---
Mandatory -->
- Did you write any new necessary tests?: N/A <!--- Mandatory for new
features or examples. -->
- Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?:
N/A <!--- Only for new features, API changes, critical bug fixes or
backward incompatible changes. -->

### Additional Information
<!-- E.g. related issue. -->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Documentation**
* Updated pruning examples and CLI help to reference directory-based
checkpoint paths for intermediate pruning scores and clarified default
wording.

* **Improvements**
* Checkpoint handling now accepts directory paths for intermediate score
storage, produces per-rank files automatically, and resolves to a
directory default when unset.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

examples/megatron_bridge/prune_minitron.py

@@ -111,8 +111,8 @@ def get_args() -> argparse.Namespace:
         type=str,
         default=None,
         help=(
-            "Path to save/restore intermediate pruning scores for resuming / faster re-run. "
-            "If not provided, it will default to `<output_path>/modelopt_pruning_scores.pth`"
+            "Directory to save/restore per-rank intermediate pruning scores for resuming / faster re-run. "
+            "If not provided, it will default to `<output_path>/modelopt_pruning_scores`"
         ),
     )
 
@@ -187,13 +187,11 @@ def get_args() -> argparse.Namespace:
     # Post-process arguments
     if args.prune_intermediate_ckpt is None:
         if args.output_megatron_path:
-            args.prune_intermediate_ckpt = (
-                f"{args.output_megatron_path}/modelopt_pruning_scores.pth"
-            )
+            args.prune_intermediate_ckpt = f"{args.output_megatron_path}/modelopt_pruning_scores"
         elif args.output_hf_path:
-            args.prune_intermediate_ckpt = f"{args.output_hf_path}/modelopt_pruning_scores.pth"
+            args.prune_intermediate_ckpt = f"{args.output_hf_path}/modelopt_pruning_scores"
         print_rank_0(
-            "No checkpoint provided to cache intermediate pruning scores. "
+            "No directory provided to cache per-rank intermediate pruning scores. "
             f"Setting to: {args.prune_intermediate_ckpt}"
         )
 

examples/pruning/README.md

@@ -97,7 +97,7 @@ This mode can be useful when you know the exact dimensions you want to prune to
 # Specify the pruning constraints (Check Support Matrix for available pruning dimensions)
 # Save minitron scores at checkpoint so we can re-run pruning with different constraints without running the forward loop again
 constraints = {"export_config": {"num_layers": 32, "hidden_size": 3584, "ffn_hidden_size": 10240}}
-config = {"forward_loop": forward_loop, "checkpoint": "/path/to/cache/pruning/scores.pth"}
+config = {"forward_loop": forward_loop, "checkpoint": "/path/to/cache/pruning/scores/"}
 
 mtp.prune(...)
 ```
@@ -129,7 +129,7 @@ def score_func(m):
 constraints = {"params": 6e9}  # Prune to 6B parameters
 config = {
     "forward_loop": forward_loop,
-    "checkpoint": "/path/to/cache/pruning/scores.pth",
+    "checkpoint": "/path/to/cache/pruning/scores/",
     "score_func": score_func,
     # Optional: Configure search space constraints (showing defaults)
     "max_width_pruning": 0.4,  # Maximum 40% per width pruning hparams (hidden_size, ffn_hidden_size, etc.)

modelopt/torch/opt/searcher.py

@@ -26,7 +26,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from contextlib import nullcontext
-from typing import Any, final
+from typing import TYPE_CHECKING, Any, final
 
 import numpy as np
 import pulp
@@ -36,6 +36,9 @@
 from modelopt.torch.utils import distributed as dist
 from modelopt.torch.utils import no_stdout, print_rank_0, run_forward_loop, warn_rank_0
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 LimitsTuple = tuple[float, float]
 ConstraintsDict = dict[str, str | float | dict | None]
 Deployment = dict[str, str]
@@ -238,9 +241,18 @@ def state_dict(self) -> SearchStateDict:
 
     def _get_checkpoint_path(self) -> str | None:
         """Get per-rank checkpoint path when distributed, otherwise the original path."""
-        checkpoint = self.config["checkpoint"]
+        checkpoint: str | Path | None = self.config["checkpoint"]
         if checkpoint is None:
             return None
+        checkpoint = str(checkpoint)
+        # Detect directory: exists as dir, ends with separator, or has no file extension
+        is_dir_path = (
+            os.path.isdir(checkpoint)
+            or checkpoint.endswith(os.sep)
+            or not os.path.splitext(checkpoint)[1]
+        )
+        if is_dir_path:
+            return os.path.join(checkpoint, f"rank{dist.rank()}.pth")
         if dist.is_initialized():
             dirname, basename = os.path.split(checkpoint)
             name, ext = os.path.splitext(basename)

tests/gpu_megatron/torch/prune/plugins/test_mcore_gpt_minitron_pruning.py

@@ -123,7 +123,7 @@ def _test_mcore_gpt_pruning(
     uneven_pp,
     position_embedding_type,
     skip_sorting,
-    ckpt_path,
+    ckpt_dir,
     rank,
     size,
 ):
@@ -196,11 +196,11 @@ def forward_loop(m):
     constraints = {"export_config": export_config}
 
     config = {
-        "checkpoint": ckpt_path,
+        "checkpoint": ckpt_dir,
         "skip_sorting": skip_sorting,
     }
     if skip_sorting:
-        assert ckpt_path is None
+        assert ckpt_dir is None
     else:
         config["forward_loop"] = forward_loop
     model, pruning_scores = prune_minitron(model, constraints, config, channel_divisor)
@@ -236,11 +236,11 @@ def forward_loop(m):
     output = run_mcore_inference(model, prompt_tokens, pruned_hidden_size)
 
     # Assert re-pruning from checkpoint works without running the forward loop again
-    if ckpt_path:
+    if ckpt_dir:
         model_rerun = _get_model(initialize_megatron=False)
         model_rerun.load_state_dict(sd)
         model_rerun, pruning_scores = prune_minitron(
-            model_rerun, constraints, {"checkpoint": ckpt_path}, channel_divisor
+            model_rerun, constraints, {"checkpoint": ckpt_dir}, channel_divisor
         )
 
         output_rerun = run_mcore_inference(model_rerun, prompt_tokens, pruned_hidden_size)
@@ -305,7 +305,7 @@ def test_mcore_gpt_pruning(
             uneven_pp,
             position_embedding_type,
             skip_sorting,
-            tmp_path / "minitron_scores.pth" if test_ckpt else None,
+            tmp_path / "minitron_scores" if test_ckpt else None,
         ),
     )
 
@@ -391,7 +391,7 @@ def test_mcore_gpt_moe_parameter_sorting(dist_workers):
     dist_workers.run(_test_mcore_gpt_moe_parameter_sorting)
 
 
-def _test_mcore_gpt_pruning_moe(ckpt_path, rank, size):
+def _test_mcore_gpt_pruning_moe(ckpt_dir, rank, size):
     channel_divisor = 4
 
     num_layers = size
@@ -442,7 +442,7 @@ def forward_loop(m):
     prune_minitron(
         model,
         constraints,
-        {"checkpoint": ckpt_path, "forward_loop": forward_loop},
+        {"checkpoint": ckpt_dir, "forward_loop": forward_loop},
         channel_divisor,
     )
 
@@ -479,14 +479,14 @@ def forward_loop(m):
     # Assert re-pruning from checkpoint works without running the forward loop again
     model_rerun = _get_model(initialize_megatron=False)
     model_rerun.load_state_dict(sd)
-    prune_minitron(model_rerun, constraints, {"checkpoint": ckpt_path}, channel_divisor)
+    prune_minitron(model_rerun, constraints, {"checkpoint": ckpt_dir}, channel_divisor)
 
     output_rerun = run_mcore_inference(model_rerun, prompt_tokens, pruned_hidden_size)
     assert torch.allclose(output, output_rerun, atol=1e-5)
 
 
 def test_mcore_gpt_pruning_moe(dist_workers, tmp_path):
-    dist_workers.run(partial(_test_mcore_gpt_pruning_moe, tmp_path / "minitron_scores.pth"))
+    dist_workers.run(partial(_test_mcore_gpt_pruning_moe, tmp_path / "minitron_scores"))
 
 
 def test_generate_search_space_combos():

tests/gpu_megatron/torch/prune/plugins/test_mcore_mamba_minitron_pruning.py

@@ -119,7 +119,7 @@ def test_mcore_mamba_parameter_sorting(dist_workers):
     dist_workers.run(_test_mcore_mamba_parameter_sorting)
 
 
-def _test_mcore_mamba_hybrid_pruning(ckpt_path, rank, size):
+def _test_mcore_mamba_hybrid_pruning(ckpt_dir, rank, size):
     channel_divisor = 4
 
     num_layers = min(size * 2, 8)
@@ -193,7 +193,7 @@ def forward_loop(m):
     prune_minitron(
         model,
         constraints,
-        {"forward_loop": forward_loop, "checkpoint": ckpt_path},
+        {"forward_loop": forward_loop, "checkpoint": ckpt_dir},
         channel_divisor,
     )
 
@@ -224,16 +224,14 @@ def forward_loop(m):
 
     # Assert re-pruning from checkpoint works without running the forward loop again
     model = _get_model(initialize_megatron=False)
-    prune_minitron(model, constraints, {"checkpoint": ckpt_path}, channel_divisor)
+    prune_minitron(model, constraints, {"checkpoint": ckpt_dir}, channel_divisor)
 
 
 def test_mcore_mamba_hybrid_pruning(dist_workers, tmp_path):
-    dist_workers.run(
-        partial(_test_mcore_mamba_hybrid_pruning, tmp_path / "modelopt_minitron_scores.pth")
-    )
+    dist_workers.run(partial(_test_mcore_mamba_hybrid_pruning, tmp_path / "minitron_scores"))
 
 
-def _test_mcore_mamba_hybrid_pruning_nas(ckpt_path, rank, size):
+def _test_mcore_mamba_hybrid_pruning_nas(ckpt_dir, rank, size):
     set_seed(SEED)
     channel_divisor = 4
 
@@ -297,7 +295,7 @@ def score_func(m):
     constraints = {"params": int(param_count * 0.7)}
     config = {
         "forward_loop": forward_loop,
-        "checkpoint": ckpt_path,
+        "checkpoint": ckpt_dir,
         "score_func": score_func,
         "max_width_pruning": 0.5,
         "max_depth_pruning": 0.5,
@@ -363,5 +361,5 @@ def score_func(m):
 )
 def test_mcore_mamba_hybrid_pruning_nas(dist_workers, tmp_path):
     dist_workers.run(
-        partial(_test_mcore_mamba_hybrid_pruning_nas, tmp_path / "modelopt_minitron_scores.pth"),
+        partial(_test_mcore_mamba_hybrid_pruning_nas, tmp_path / "minitron_scores"),
     )