Merge branch 'main' into shengliangx/docs-preview

Author: shengliangxu

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

@@ -98,7 +98,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(...)
 ```
@@ -130,7 +130,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

@@ -124,7 +124,7 @@ def _test_mcore_gpt_pruning(
     uneven_pp,
     position_embedding_type,
     skip_sorting,
-    ckpt_path,
+    ckpt_dir,
     rank,
     size,
 ):
@@ -198,11 +198,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)
@@ -238,11 +238,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)
@@ -307,7 +307,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,
         ),
     )
 
@@ -394,7 +394,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
@@ -446,7 +446,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,
     )
 
@@ -483,14 +483,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

@@ -120,7 +120,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)
@@ -196,7 +196,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
 
@@ -299,7 +297,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,
@@ -365,5 +363,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"),
     )