Add scoring module

Signed-off-by: Daniel Korzekwa <[email protected]>

Author: danielkorzekwa

modelopt/torch/_compress/nas/plugins/compress_nas_plugin.py

@@ -21,12 +21,12 @@
 from pathlib import Path
 
 import mip_and_realize_models
-import scoring
 import torch
 from torch import nn
 
 import modelopt.torch._compress.build_library_and_stats as build_library_and_stats
 import modelopt.torch._compress.pruning.pruning_ckpts as pruning_ckpts
+import modelopt.torch._compress.scoring.scoring as scoring
 from modelopt.torch._compress.activation_scoring import score_pruning_activations
 from modelopt.torch._compress.decilm.converters.convert_llama3_to_decilm import (
     convert_llama3_to_decilm,

modelopt/torch/_compress/scoring/scoring.py

@@ -0,0 +1,97 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# mypy: ignore-errors
+import os
+import re
+from glob import glob
+from pathlib import Path
+
+import hydra
+import numpy as np
+import pandas as pd
+import torch
+from omegaconf import DictConfig
+
+from modelopt.torch._compress.tools.hydra_utils import register_hydra_resolvers
+from modelopt.torch._compress.tools.logger import mprint
+from modelopt.torch._compress.tools.runtime import BaseRuntime, IRuntime, NativeDdpRuntime
+from modelopt.torch._compress.tools.validate_puzzle_with_multi_replacements import (
+    validate_puzzle_solutions,
+)
+from modelopt.torch._compress.utils.dist_utils import is_distributed
+
+
+def extract_solution_id(filename):
+    pattern = r"solution_(\d+)\.json"
+    match = re.search(pattern, filename)
+
+    if match:
+        solution_id = match.group(1)
+        return int(solution_id)
+    else:
+        mprint(f"Couldn't extract solutions_id from file {filename}")
+
+
+def find_missing_solutions(solutions_df, validation_dir):
+    all_solutions = np.arange(solutions_df.shape[0])
+
+    benchmarked_solutions = list(glob(f"{validation_dir}/solution*.json"))
+    benchmarked_solutions = [
+        extract_solution_id(os.path.basename(s)) for s in benchmarked_solutions
+    ]
+    benchmarked_solutions = [s for s in benchmarked_solutions if s is not None]
+
+    unbenchmarked_solutions = np.setdiff1d(all_solutions, benchmarked_solutions)
+    return unbenchmarked_solutions.tolist()
+
+
+def get_solutions_to_validate(cfg: DictConfig):
+    _solutions_to_validate = cfg.scoring.solutions_to_validate
+    if _solutions_to_validate is None:
+        single_block_replacement_solutions = pd.read_json(cfg.scoring.solutions_path)
+        if cfg.scoring.skip_existing_solutions:
+            _solutions_to_validate = find_missing_solutions(
+                single_block_replacement_solutions, cfg.scoring.output_dir
+            )
+        else:
+            _solutions_to_validate = np.arange(single_block_replacement_solutions.shape[0]).tolist()
+    return _solutions_to_validate
+
+
+def launch_scoring(cfg: DictConfig, runtime: IRuntime):
+    cfg.scoring.solutions_to_validate = get_solutions_to_validate(cfg)
+    mprint(f"Solutions to validate: {cfg.scoring.solutions_to_validate}")
+    validate_puzzle_solutions(args=cfg.scoring, runtime=runtime)
+
+
+@hydra.main("", version_base="1.3")
+def main(cfg: DictConfig) -> None:
+    cfg = hydra.utils.instantiate(cfg)
+    mprint(cfg)
+
+    _runtime = (
+        NativeDdpRuntime(
+            dtype=torch.bfloat16, torch_distributed_timeout=getattr(cfg, "nccl_timeout_minutes")
+        )
+        if is_distributed()
+        else BaseRuntime(dtype=torch.bfloat16)
+    )
+    with _runtime as runtime:
+        launch_scoring(cfg, runtime)
+
+
+if __name__ == "__main__":
+    register_hydra_resolvers()
+    main()