Add score calculation logic

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

Author: kevalmorabia97

modelopt/torch/prune/plugins/mcore_minitron.py

@@ -25,6 +25,7 @@
 """
 
 from collections.abc import Callable
+from dataclasses import dataclass
 from functools import partial
 from itertools import product
 from typing import Any
@@ -169,6 +170,13 @@ def drop_mcore_language_model_layers(model: nn.Module, *, layers_to_drop: list[i
     model.config.num_layers = new_num_layers
 
 
+@dataclass
+class CandidateSubnet:
+    ss_config: dict
+    params: float
+    score: float | None
+
+
 class MCoreMinitronSearcher(BaseSearcher):
     """Searcher for Minitron pruning algorithm.
 
@@ -182,7 +190,8 @@ class MCoreMinitronSearcher(BaseSearcher):
 
     activations_per_rank: list[dict[str, torch.Tensor]]
     layer_scores: dict[int, torch.Tensor]
-    top_k_candidates_per_constraint: dict[float, list[tuple[dict, float]]]
+    # Dict from params constraint to list of tuples (ss_config, params, score)
+    top_k_candidates_per_constraint: dict[float, list[CandidateSubnet]]
 
     @property
     def default_search_config(self) -> SearchConfig:
@@ -400,7 +409,7 @@ def search_best_arch_by_params(self, sorted_layers: list[int]) -> dict:
                 max_depth_pruning,
             )
             sample(self.model, sample_func=max)  # reset to max subnet (for sanity)
-            selected: list[tuple[dict, float]] = []
+            selected = []
             for ss_config in tqdm(
                 search_space_configs,
                 desc=f"Finding top {top_k} candidates fitting the constraints...",
@@ -415,23 +424,40 @@ def search_best_arch_by_params(self, sorted_layers: list[int]) -> dict:
                     layer_ids = sorted_layers[: ss_config["num_layers"]]
                 candidate_params = _param_num_dynamic(self.model, layer_numbers_to_count=layer_ids)
                 if candidate_params <= max_params:
-                    selected.append((ss_config, candidate_params))
+                    selected.append(CandidateSubnet(ss_config, candidate_params, None))
                 sample(self.model, sample_func=max)  # reset to max subnet
             assert len(selected) > 0, "No subnets found fitting the constraints!"
             print_rank_0(f"Found {len(selected)} candidates fitting the constraints!")
             self.top_k_candidates_per_constraint[max_params] = sorted(
-                selected, key=lambda x: x[1], reverse=True
+                selected, key=lambda x: x.params, reverse=True
             )[:top_k]
             self.save_search_checkpoint(verbose=True)
         else:
             print_rank_0(f"Using top {top_k} candidates from checkpoint")
         top_k_candidates = self.top_k_candidates_per_constraint[max_params]
 
+        print_rank_0(f"\n====================\nTop {top_k} candidates:")
+        for candidate in top_k_candidates:
+            print_rank_0(f"\t{num2hrb(candidate.params)} params: {candidate.ss_config}")
+        print_rank_0("====================\n")
+
         # 3. Validate top-k candidates using the score_func and return the best subnet
-        # TODO: update this
-        best = top_k_candidates[0][0]
+        for candidate in tqdm(
+            top_k_candidates,
+            desc=f"Validating top {top_k} candidates on given score_func...",
+            disable=not dist.is_master(),
+        ):
+            if candidate.score is None:  # not restored from checkpoint
+                self._prune(candidate.ss_config, prune_depth=False, update_config=False)
+                candidate.score = self.eval_score(silent=True)
+                sample(self.model, sample_func=max)  # reset to max subnet
+                self.save_search_checkpoint(verbose=False)
+            print_rank_0(
+                f"\t{num2hrb(candidate.params)} params: {candidate.ss_config} -> {candidate.score:.4f} score"
+            )
 
-        return best
+        best = max(top_k_candidates, key=lambda x: x.score)  # type: ignore[arg-type, return-value]
+        return best.ss_config
 
     @staticmethod
     def _generate_search_space_combos(