From ed5c00f75a8bdebc73a4920bab82f5c925f5132c Mon Sep 17 00:00:00 2001 From: Daniel Korzekwa Date: Fri, 6 Mar 2026 06:11:59 -0800 Subject: [PATCH] merge any_model: mip_and_realize_models Signed-off-by: Daniel Korzekwa --- modelopt/torch/puzzletron/mip/run_puzzle.py | 4 +- modelopt/torch/puzzletron/puzzletron.py | 4 +- tests/gpu/torch/puzzletron/test_puzzletron.py | 108 +++++++++--------- 3 files changed, 57 insertions(+), 59 deletions(-) diff --git a/modelopt/torch/puzzletron/mip/run_puzzle.py b/modelopt/torch/puzzletron/mip/run_puzzle.py index 72919d27c..da0f90452 100644 --- a/modelopt/torch/puzzletron/mip/run_puzzle.py +++ b/modelopt/torch/puzzletron/mip/run_puzzle.py @@ -688,9 +688,7 @@ def _get_block_stats( not (block_config.attention.no_op and block_config.ffn.no_op) ) block_stats["num_kv_heads"] = ( - subblock_stats["args"]["n_head"] // block_config.attention.n_heads_in_group - if block_stats["has_attention"] - else 0 + block_config.attention.num_key_value_heads if block_stats["has_attention"] else 0 ) block_stats["num_local_experts"] = ( block_config.ffn.moe.num_local_experts if block_stats["has_moe"] else 0 diff --git a/modelopt/torch/puzzletron/puzzletron.py b/modelopt/torch/puzzletron/puzzletron.py index 262df7648..5a1484e07 100644 --- a/modelopt/torch/puzzletron/puzzletron.py +++ b/modelopt/torch/puzzletron/puzzletron.py @@ -70,7 +70,7 @@ def puzzletron( # Step 5: calc_one_block_scores (distributed processing) scoring.launch_scoring(hydra_cfg) - # # Step 6: mip_and_realize_models (distributed processing) - # mip_and_realize_models.launch_mip_and_realize_model(hydra_cfg) + # Step 6: mip_and_realize_models (distributed processing) + mip_and_realize_models.launch_mip_and_realize_model(hydra_cfg) return hydra_cfg diff --git a/tests/gpu/torch/puzzletron/test_puzzletron.py b/tests/gpu/torch/puzzletron/test_puzzletron.py index 585567715..fbaaf85a1 100644 --- a/tests/gpu/torch/puzzletron/test_puzzletron.py +++ b/tests/gpu/torch/puzzletron/test_puzzletron.py @@ -126,60 +126,60 @@ def _test_puzzletron_multiprocess_job( str(hydra_config_dir), hydra_config_subdir, str(puzzle_dir), str(dataset_path) ) - # # - # # Check assertions - # # - # if rank == 0: - # if has_moe_layers: - # # assertions for the score_pruning_activations step 1 (MoE models only) - # rank_filepath = ( - # f"pruning/pruning_scores/expert_removal/10samples_diverse_mini/rank_{rank}.pth" - # ) - # assert (puzzle_dir / rank_filepath).is_file(), f"Expected {rank_filepath} to exist" - - # # assertions for the pruning_ckpts step 2 - # assert (puzzle_dir / "ckpts/num_experts_8").exists() - - # # assertions for the mip_and_realize_models step 6 - # # Find the MIP solution directory dynamically (e.g., stats_num_local_experts_*) - # mip_solutions_dir = puzzle_dir / "mip/puzzle_solutions" - # solution_dirs = [ - # d - # for d in mip_solutions_dir.iterdir() - # if d.is_dir() and d.name.startswith("stats_num_local_experts_") - # ] - # assert len(solution_dirs) == 1, ( - # f"Expected exactly one stats_num_local_experts_* directory, found: {[d.name for d in solution_dirs]}" - # ) - # solution_dir = solution_dirs[0] - - # solution_0_ckpt_config_path = ( - # solution_dir / "solutions--checkpoints/solution_0/config.json" - # ) - # assert solution_0_ckpt_config_path.exists() - # assert (solution_dir / "solutions.json").exists() - - # # Validate lm_loss - # _assert_lm_loss(puzzle_dir, hf_config_name) - # else: - # # assertions for the score_pruning_activations step 1 (FFN pruning) - # _assert_score_pruning_activations(puzzle_dir, hf_config_name) - - # # assertions for the pruning_ckpts step 2 - # assert (puzzle_dir / "ckpts/ffn_256_attn_no_op").exists() - - # # assertions for the mip_and_realize_models step 6 - # _assert_mip_solutions(puzzle_dir, hf_config_name) - - # # assertions for the build_library_and_stats step 4 - # assert (puzzle_dir / "replacement_library.json").is_file() - # assert (puzzle_dir / "subblock_stats.json").is_file() - - # # assertions for the scoring step 5 - # solution_0_filepath = ( - # puzzle_dir / "single_sequence_replacement_solutions--validation/solution_0.json" - # ) - # assert solution_0_filepath.exists() + # + # Check assertions + # + if rank == 0: + if has_moe_layers: + # assertions for the score_pruning_activations step 1 (MoE models only) + rank_filepath = ( + f"pruning/pruning_scores/expert_removal/10samples_diverse_mini/rank_{rank}.pth" + ) + assert (puzzle_dir / rank_filepath).is_file(), f"Expected {rank_filepath} to exist" + + # assertions for the pruning_ckpts step 2 + assert (puzzle_dir / "ckpts/num_experts_8").exists() + + # assertions for the mip_and_realize_models step 6 + # Find the MIP solution directory dynamically (e.g., stats_num_local_experts_*) + mip_solutions_dir = puzzle_dir / "mip/puzzle_solutions" + solution_dirs = [ + d + for d in mip_solutions_dir.iterdir() + if d.is_dir() and d.name.startswith("stats_num_local_experts_") + ] + assert len(solution_dirs) == 1, ( + f"Expected exactly one stats_num_local_experts_* directory, found: {[d.name for d in solution_dirs]}" + ) + solution_dir = solution_dirs[0] + + solution_0_ckpt_config_path = ( + solution_dir / "solutions--checkpoints/solution_0/config.json" + ) + assert solution_0_ckpt_config_path.exists() + assert (solution_dir / "solutions.json").exists() + + # Validate lm_loss + _assert_lm_loss(puzzle_dir, hf_config_name) + else: + # assertions for the score_pruning_activations step 1 (FFN pruning) + _assert_score_pruning_activations(puzzle_dir, hf_config_name) + + # assertions for the pruning_ckpts step 2 + assert (puzzle_dir / "ckpts/ffn_256_attn_no_op").exists() + + # assertions for the mip_and_realize_models step 6 + _assert_mip_solutions(puzzle_dir, hf_config_name) + + # assertions for the build_library_and_stats step 4 + assert (puzzle_dir / "replacement_library.json").is_file() + assert (puzzle_dir / "subblock_stats.json").is_file() + + # assertions for the scoring step 5 + solution_0_filepath = ( + puzzle_dir / "single_sequence_replacement_solutions--validation/solution_0.json" + ) + assert solution_0_filepath.exists() dist.cleanup()