From a724afd0ee38c7a62f8206bf2cc11dcaa4fe8685 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Mon, 15 Jun 2026 12:09:27 -0700 Subject: [PATCH 1/8] adding some fixes for puzzletron/runtime tutorial Signed-off-by: Grzegorz Karch --- .../Llama-3_1-8B.yaml | 8 ++--- .../pruning/attn_pruning.yaml | 23 +++++++++++++ .../pruning/ffn_pruning.yaml | 19 +++++++++++ .../pruning/hidden_dim_pruning.yaml | 15 +++++++++ .../pruning/pruning_defaults.yaml | 33 +++++++++++++++++++ .../validate_model_defaults.yaml | 17 ++++++++++ .../validate_solutions_defaults.yaml | 10 ++++++ .../subblock_stats/calc_runtime_stats.py | 1 + .../subblock_stats/runtime_utils.py | 32 ++++++++++++++++++ .../puzzletron/subblock_stats/runtime_vllm.py | 8 ++++- 10 files changed, 161 insertions(+), 5 deletions(-) create mode 100644 examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/attn_pruning.yaml create mode 100644 examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/ffn_pruning.yaml create mode 100644 examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/hidden_dim_pruning.yaml create mode 100644 examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/pruning_defaults.yaml create mode 100644 examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml create mode 100644 examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_solutions_defaults.yaml diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/Llama-3_1-8B.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/Llama-3_1-8B.yaml index b4adbb82add..74aa609f44d 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/Llama-3_1-8B.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/Llama-3_1-8B.yaml @@ -1,7 +1,7 @@ defaults: - - ../llama-3_1-8B_pruneffn_memory/pruning/ffn_pruning@pruning - - ../llama-3_1-8B_pruneffn_memory/validate_solutions_defaults@scoring - - ../llama-3_1-8B_pruneffn_memory/validate_solutions_defaults@realize_model + - pruning: ffn_pruning + - scoring: ../validate_solutions_defaults + - realize_model: ../validate_solutions_defaults - bypass: - override hydra/hydra_logging: disabled - _self_ @@ -39,7 +39,7 @@ scoring: teacher_dir: ${to_path:${teacher_dir}} output_dir: ${puzzle_dir}/single_sequence_replacement_solutions--validation - eval_samples: 128 + eval_samples: 16 micro_batch_size: 1 seed: 42 shuffle_seed: 444 diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/attn_pruning.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/attn_pruning.yaml new file mode 100644 index 00000000000..53d7e4bd9c6 --- /dev/null +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/attn_pruning.yaml @@ -0,0 +1,23 @@ +defaults: + - pruning_defaults + +hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IndependentKvHeadContributionHook} + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/attn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +pruning_mixin: + _target_: modelopt.torch.puzzletron.pruning.kv_heads_pruning_mixin.KVHeadsPruningMixIn + layer_descriptor: + _target_: modelopt.torch.puzzletron.anymodel.models.llama.llama_model_descriptor.LlamaKVHeadsLayerDescriptor + +activation_hooks_kwargs: + method: independent_kv_head_contribution + optimize_for: memory # IndependentKvHeadContributionHook implementation that consumes less memory + target_layer: "self_attn.o_proj" + layer_input_descriptors_path: + +# n_heads_in_group: 4 +# num_attention_heads: 32 # num query heads +# num_kv_heads: 32 / 4 = 8 # num_query_heads // n_heads_in_group +n_heads_in_group_list: [8, 16, 32] # num_kv_heads = [4, 2, 1] +gqa_init_mode: "PruneKVHeads" diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/ffn_pruning.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/ffn_pruning.yaml new file mode 100644 index 00000000000..da0b9720700 --- /dev/null +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/ffn_pruning.yaml @@ -0,0 +1,19 @@ +defaults: + - pruning_defaults + +pruning_mixin: + _target_: modelopt.torch.puzzletron.pruning.ffn_intermediate_pruning_mixin.FFNIntermediatePruningMixIn + layer_descriptor: + _target_: modelopt.torch.puzzletron.anymodel.models.llama.llama_model_descriptor.LlamaFFNIntermediateLayerDescriptor + +hook_class: ${get_object:modelopt.torch.prune.importance_hooks.base_hooks.IterativeChannelContributionHook} + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/ffn_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: iterative + target_layer: "mlp.down_proj" + layer_input_descriptors_path: + +intermediate_size_list: [3072, 5888, 8704, 11520] # teacher_intermediate_size is 14336 +mlp_init_mode: "PruneByActivationsLog" diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/hidden_dim_pruning.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/hidden_dim_pruning.yaml new file mode 100644 index 00000000000..407c835d8c4 --- /dev/null +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/hidden_dim_pruning.yaml @@ -0,0 +1,15 @@ +defaults: + - pruning_defaults + +activations_log_dir: ${puzzle_dir}/pruning/pruning_scores/hidden_dim_${pruning.activation_hooks_kwargs.method}/${pruning.experiment_id} + +activation_hooks_kwargs: + method: layer_norm_contribution + target_layer: "layernorm" + +# Hidden dimension pruning specific settings +hidden_size_list: [3072, 2048] # Target hidden sizes to prune to +hidden_size_init_mode: "PruneByChannelRanking" +mlp_init_mode: "Truncate" # TODO, make it work with CopyAsIs/FromTeacher +gqa_init_mode: "AverageKV" # TODO, make it work with CopyAsIs/FromTeacher +linear_init_mode: "FromTeacher" diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/pruning_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/pruning_defaults.yaml new file mode 100644 index 00000000000..e05e775bee3 --- /dev/null +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/pruning/pruning_defaults.yaml @@ -0,0 +1,33 @@ +defaults: + - /validate_model_defaults + +descriptor: ${descriptor} +model_name_or_path: ${teacher_dir} +experiment_id: ${pruning.eval_samples}samples_diverse_mini +activations_log_dir: ??? +activation_hooks_kwargs: ??? + +# Data: +eval_samples: 1000 # default is 10000 +micro_batch_size: 4 +dataset_path: ${dataset_path} +val_dataset_name: train + +# Prune ckpts +pruned_ckpts_output_dir: ${puzzle_dir}/pruning/${pruning.experiment_id} + +## FFN pruning +ffn_list: +mlp_init_mode: "Truncate" # PruneByActivationsLog + +## KV-heads pruning +n_heads_in_group_list: +gqa_init_mode: "AverageKV" + +## Hidden dimension pruning +hidden_size_list: +hidden_size_init_mode: "PruneByChannelRanking" +linear_init_mode: "FromTeacher" + +mlp_init_config_yaml: + activations_log_dir: ${pruning.activations_log_dir} diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml new file mode 100644 index 00000000000..ce1749d9698 --- /dev/null +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml @@ -0,0 +1,17 @@ +model_dtype: torch.bfloat16 # dtype to cast the model for validate_model +autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model +block_size: 8192 +bos_rate: 0.5 +data_column: messages +val_dataset_name: valid +shuffle_seed: 81436 +seed: 42 +fim_rate: 0 +fim_spm_rate: 0 +source_datasets_to_discard: +varlen: false +write_results: false +calc_losses_on_cpu: false +activations_log_dir: +model_name_or_path: +load_dataset_fn: ${get_object:modelopt.torch.puzzletron.utils.data.dataloaders.load_from_disk_fn} diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_solutions_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_solutions_defaults.yaml new file mode 100644 index 00000000000..ec139023794 --- /dev/null +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_solutions_defaults.yaml @@ -0,0 +1,10 @@ +defaults: + - /validate_model_defaults + - _self_ + +solutions_to_validate: +skip_validation: false +save_models: false +bigger_is_better: false +sort_solutions_by: +calculate_full_score_ablations: false diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py index 6e4821936e7..c7d926de18f 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py @@ -188,6 +188,7 @@ def calc_runtime_for_subblocks( batch_size, runtime_stats_config.get("num_iters", 30), runtime_stats_config.get("num_warmup_iters", 10), + runtime_stats_config.get("gpu_memory_utilization", 0.5), ) runtime_by_subblock_dict = {} diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index 3259e706c73..7ea8fabdf03 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -25,12 +25,15 @@ import json from dataclasses import dataclass from pathlib import Path +from types import SimpleNamespace import torch from transformers import AutoTokenizer, LlamaForCausalLM from ..anymodel.converter import Converter from ..anymodel.models.llama import LlamaModelDescriptor +from ..tools.logger import mprint +from ..utils.vllm_adapter import convert_block_configs_to_per_layer_config @dataclass(frozen=True) @@ -48,6 +51,11 @@ class RuntimeConfig: batch_size: int num_iters: int num_warmup_iters: int + # Fraction of total GPU memory vLLM may use. Kept well below the default + # (~0.9) because the parent puzzletron process is co-resident on the same + # GPU during benchmarking; requesting too much fails vLLM's startup + # free-memory check. + gpu_memory_utilization: float = 0.5 def save_model(model: LlamaForCausalLM, tokenizer_path: Path, output_path: Path) -> None: @@ -80,3 +88,27 @@ def save_model_as_anymodel(model, output_dir: Path, descriptor): config_data["architectures"] = ["AnyModel"] with open(config_path, "w") as f: json.dump(config_data, f, indent=2) + + +def convert_config_to_vllm_anymodel(input_dir: Path, output_dir: Path): + """Convert a model to vLLM AnyModel format.""" + # Load the model config.json, update "architectures" to ["AnyModel"], and write back to disk. + input_config_path = Path(input_dir) / "config.json" + if not input_config_path.exists(): + raise FileNotFoundError(f"Config file not found at {input_config_path}") + try: + with open(input_config_path) as f: + config_data = json.load(f) + except json.JSONDecodeError as e: + raise ValueError(f"Error loading config file: {e}") from e + + config = SimpleNamespace(**config_data) + config.architectures = ["AnyModel"] + config.base_architecture = "LlamaForCausalLM" + + if convert_block_configs_to_per_layer_config(config): + mprint("Converted block configs to per-layer config") + else: + mprint("No block configs to convert") + with open(Path(output_dir) / "config.json", "w") as f: + json.dump(vars(config), f, indent=2) diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py index 14eb337b707..386f3615d49 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py @@ -29,6 +29,7 @@ import json import subprocess # nosec B404 from pathlib import Path +from types import SimpleNamespace from ..tools.logger import mprint from ..utils.vllm_adapter import convert_block_configs_to_per_layer_config @@ -48,10 +49,11 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) with open(model_path / "config.json") as f: config = json.load(f) + config = SimpleNamespace(**config) if convert_block_configs_to_per_layer_config(config): mprint("Converted block configs to per-layer config") with open(model_path / "config.json", "w") as f: - json.dump(config, f, indent=2) + json.dump(vars(config), f, indent=2) else: mprint("No block configs to convert") @@ -83,6 +85,10 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) "1", "--distributed-executor-backend", "external_launcher", + # Cap GPU memory so vLLM's startup free-memory check passes while the + # parent puzzletron process is co-resident on the same GPU. + "--gpu-memory-utilization", + str(runtime_config.gpu_memory_utilization), # Required for accurate per-block runtime stats. "--optimization-level", "0", From ff55eee56ee1e185f8120a2c30b1deee8661e61a Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Tue, 23 Jun 2026 02:50:14 -0700 Subject: [PATCH 2/8] validation->valid in llama 3.1 config Signed-off-by: Grzegorz Karch --- .../llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml index 6b36142a3a8..ce1749d9698 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: validation +val_dataset_name: valid shuffle_seed: 81436 seed: 42 fim_rate: 0 From b86f13c7cd191407dd3b9b34d89606e0c61198d7 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Thu, 25 Jun 2026 13:09:19 -0700 Subject: [PATCH 3/8] reverting unnecessary change Signed-off-by: Grzegorz Karch --- .../llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 From 0688d00d1c507bc50f8adc3d3137e7db6149d06d Mon Sep 17 00:00:00 2001 From: "Grzegorz K. Karch" Date: Thu, 2 Jul 2026 17:07:32 +0200 Subject: [PATCH 4/8] Rename validation dataset from 'valid' to 'validation' Signed-off-by: Grzegorz K. Karch --- .../llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml index ce1749d9698..6b36142a3a8 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/validate_model_defaults.yaml @@ -3,7 +3,7 @@ autocast_dtype: torch.bfloat16 # dtype for torch.autocast for validate_model block_size: 8192 bos_rate: 0.5 data_column: messages -val_dataset_name: valid +val_dataset_name: validation shuffle_seed: 81436 seed: 42 fim_rate: 0 From bfb36198c028cb6a2dccc273d2dfad81b70d490b Mon Sep 17 00:00:00 2001 From: "Grzegorz K. Karch" Date: Thu, 2 Jul 2026 17:11:28 +0200 Subject: [PATCH 5/8] Add TODO for extending model support in runtime_utils Added a TODO comment to extend support for other models. Signed-off-by: Grzegorz K. Karch --- modelopt/torch/puzzletron/subblock_stats/runtime_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index 7ea8fabdf03..ac1ed508506 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -104,7 +104,7 @@ def convert_config_to_vllm_anymodel(input_dir: Path, output_dir: Path): config = SimpleNamespace(**config_data) config.architectures = ["AnyModel"] - config.base_architecture = "LlamaForCausalLM" + config.base_architecture = "LlamaForCausalLM" # TODO: extend support to other models if convert_block_configs_to_per_layer_config(config): mprint("Converted block configs to per-layer config") From 32bd535cb23452f12c89d7fa336f42c67ff55a5b Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 3 Jul 2026 01:34:02 -0700 Subject: [PATCH 6/8] ruff fix Signed-off-by: Grzegorz Karch --- modelopt/torch/puzzletron/subblock_stats/runtime_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index ac1ed508506..735c3f8f72b 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -104,7 +104,7 @@ def convert_config_to_vllm_anymodel(input_dir: Path, output_dir: Path): config = SimpleNamespace(**config_data) config.architectures = ["AnyModel"] - config.base_architecture = "LlamaForCausalLM" # TODO: extend support to other models + config.base_architecture = "LlamaForCausalLM" # TODO: extend support to other models if convert_block_configs_to_per_layer_config(config): mprint("Converted block configs to per-layer config") From faac367a1ba0c82732bc535f5124f6926b522fc9 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Wed, 8 Jul 2026 04:04:53 -0700 Subject: [PATCH 7/8] added instructions to convert model config to vllm anymodel Signed-off-by: Grzegorz Karch --- examples/puzzletron/README.md | 13 ++++------ .../subblock_stats/runtime_utils.py | 26 ++++++++++++++----- 2 files changed, 25 insertions(+), 14 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 9b42406667c..6beed07467a 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -276,17 +276,14 @@ See [vLLM documentation](https://docs.vllm.ai/en/latest/getting_started/installa **NOTE:** This is a temporary workaround pending official vLLM integration. You can track merge status [here](https://github.com/vllm-project/vllm/pull/36512). -Then, add the following to the model's `config.json` file (here we use Llama as an example): +Then, convert the model's config.json to AnyModel format: -```json -{ - ... - "architectures": ["AnyModel"], - "base_architecture": "LlamaForCausalLM", - ... -} +```bash +python -m modelopt.torch.puzzletron.subblock_stats.runtime_utils convert_config_to_vllm_anymodel ``` +This will create a backup of the original config.json file at `config.bak`. + For new architectures that are not supported by vLLM, you additionally need to add the following to the `config.json` file (using Llama3 as an example): ```json diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index 735c3f8f72b..204c2a74305 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -23,6 +23,7 @@ """ import json +import shutil from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace @@ -90,14 +91,21 @@ def save_model_as_anymodel(model, output_dir: Path, descriptor): json.dump(config_data, f, indent=2) -def convert_config_to_vllm_anymodel(input_dir: Path, output_dir: Path): +def convert_config_to_vllm_anymodel(config_dir: Path): """Convert a model to vLLM AnyModel format.""" # Load the model config.json, update "architectures" to ["AnyModel"], and write back to disk. - input_config_path = Path(input_dir) / "config.json" - if not input_config_path.exists(): - raise FileNotFoundError(f"Config file not found at {input_config_path}") + config_path = Path(config_dir) / "config.json" + if not config_path.exists(): + raise FileNotFoundError(f"Config file not found at {config_path}") + + backup_config_path = config_path.with_suffix(".bak") + if backup_config_path.exists(): + raise FileExistsError(f"Backup config file already exists at {backup_config_path}") + + shutil.copy(config_path, backup_config_path) + try: - with open(input_config_path) as f: + with open(config_path) as f: config_data = json.load(f) except json.JSONDecodeError as e: raise ValueError(f"Error loading config file: {e}") from e @@ -110,5 +118,11 @@ def convert_config_to_vllm_anymodel(input_dir: Path, output_dir: Path): mprint("Converted block configs to per-layer config") else: mprint("No block configs to convert") - with open(Path(output_dir) / "config.json", "w") as f: + with open(config_path, "w") as f: json.dump(vars(config), f, indent=2) + + +if __name__ == "__main__": + import fire + + fire.Fire() From 558cf708cc8323181337dd268b5fa62bea1965ad Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Fri, 10 Jul 2026 06:32:07 -0700 Subject: [PATCH 8/8] wip Signed-off-by: Grzegorz Karch --- .../subblock_stats/calc_runtime_stats.py | 186 ++++++++++++++++-- .../subblock_stats/calc_subblock_stats.py | 2 + .../subblock_stats/runtime_utils.py | 10 +- .../puzzletron/test_calc_runtime_stats.py | 31 ++- 4 files changed, 214 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py index c7d926de18f..f00f9c5c495 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py @@ -16,6 +16,7 @@ """Runtime statistics calculation for NAS subblock benchmarking via vLLM.""" +import copy import tempfile from dataclasses import replace from functools import cache @@ -23,11 +24,13 @@ from omegaconf import DictConfig from tqdm import tqdm -from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM +from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM, PretrainedConfig +from ..anymodel.model_descriptor import ModelDescriptor from ..anymodel.models.llama import LlamaModelDescriptor from ..anymodel.puzzformer import deci_x_patcher from ..block_config import AttentionConfig, BlockConfig, FFNConfig, SubblockConfig +from ..tools.checkpoint_utils_hf import init_model_from_config from .runtime_utils import RuntimeConfig, save_model from .runtime_vllm import run_vllm_latency_benchmark @@ -82,16 +85,110 @@ def create_benchmark_model( return model -def calc_model_runtime(model: LlamaForCausalLM, runtime_config: RuntimeConfig) -> float: +def _uses_moe(subblock_config: SubblockConfig | BlockConfig | None) -> bool: + if subblock_config is None: + return False + if isinstance(subblock_config, FFNConfig): + return subblock_config.is_moe + if isinstance(subblock_config, BlockConfig): + if subblock_config.ffn is not None and subblock_config.ffn.is_moe: + return True + if subblock_config.parallel_blocks is not None: + return any(_uses_moe(block_config) for block_config in subblock_config.parallel_blocks) + return False + + +def _block_config_to_pattern_char(block_config: BlockConfig) -> str: + if block_config.ffn is not None and block_config.ffn.is_moe: + return "E" + if block_config.attention is not None and block_config.attention.is_mamba: + return "M" + if block_config.attention is not None and not block_config.attention.no_op: + return "*" + return "-" + + +def _hybrid_override_pattern_for_block_configs(block_configs: list[BlockConfig]) -> str: + return "".join(_block_config_to_pattern_char(block_config) for block_config in block_configs) + + +def _make_moe_ffn_block_config(ffn_config: FFNConfig) -> BlockConfig: + return BlockConfig(attention=AttentionConfig(no_op=True), ffn=ffn_config) + + +def _make_moe_ffn_baseline_block_config() -> BlockConfig: + return BlockConfig(attention=AttentionConfig(no_op=True), ffn=FFNConfig(no_op=True)) + + +def create_descriptor_benchmark_model( + model_config: PretrainedConfig, + descriptor: type[ModelDescriptor], + vocab_size: int, + hidden_size: int, + num_key_value_heads: int, + num_attention_heads: int, + prefill_seq_len: int, + generation_seq_len: int, + block_config: BlockConfig, + repeat_block_n_times: int, + pattern_block_config: BlockConfig | None = None, +): + """Build a benchmark model with the caller's descriptor, preserving real MoE layers.""" + block_configs = [block_config] * repeat_block_n_times + pattern_block_configs = [pattern_block_config or block_config] * repeat_block_n_times + + benchmark_config = copy.deepcopy(model_config) + lm_config = descriptor.get_language_model_config(benchmark_config) + lm_config.num_hidden_layers = len(block_configs) + lm_config.hidden_size = hidden_size + lm_config.num_attention_heads = num_attention_heads + if hasattr(lm_config, "num_key_value_heads"): + lm_config.num_key_value_heads = num_key_value_heads + if hasattr(lm_config, "vocab_size"): + lm_config.vocab_size = vocab_size + if hasattr(benchmark_config, "vocab_size"): + benchmark_config.vocab_size = vocab_size + if hasattr(lm_config, "max_position_embeddings"): + lm_config.max_position_embeddings = prefill_seq_len + generation_seq_len + if hasattr(lm_config, "hybrid_override_pattern"): + lm_config.hybrid_override_pattern = _hybrid_override_pattern_for_block_configs( + pattern_block_configs + ) + + benchmark_config.block_configs = block_configs + if lm_config is not benchmark_config: + lm_config.block_configs = block_configs + + with deci_x_patcher(descriptor, block_configs): + model = init_model_from_config( + benchmark_config, + trust_remote_code=descriptor.requires_trust_remote_code(), + ) + + block_config_dicts = [block_config.to_dict() for block_config in block_configs] + model.config.block_configs = block_config_dicts + model_lm_config = descriptor.get_language_model_config(model.config) + if model_lm_config is not model.config: + model_lm_config.block_configs = block_config_dicts + model.config.architectures = ["AnyModel"] + model.config.base_architecture = type(model).__name__ + return model + + +def calc_model_runtime( + model: LlamaForCausalLM, + runtime_config: RuntimeConfig, + descriptor: type[ModelDescriptor] = LlamaModelDescriptor, +) -> float: """Measure total runtime of a model via vLLM latency benchmark.""" with tempfile.TemporaryDirectory() as model_tmpdir: - save_model(model, Path(runtime_config.tokenizer_path), Path(model_tmpdir)) + save_model(model, Path(runtime_config.tokenizer_path), Path(model_tmpdir), descriptor) model_total_runtime_ms = run_vllm_latency_benchmark(Path(model_tmpdir), runtime_config) return model_total_runtime_ms @cache -def calc_subblock_runtime( +def _calc_llama_subblock_runtime( runtime_config: RuntimeConfig, subblock_config: SubblockConfig | None, ) -> float: @@ -128,22 +225,83 @@ def calc_subblock_runtime( @cache -def calc_base_runtime(runtime_config: RuntimeConfig, subblock_config: SubblockConfig) -> float: +def _calc_llama_base_runtime( + runtime_config: RuntimeConfig, subblock_config: SubblockConfig +) -> float: """Calculate the base runtime of a model with no subblocks.""" base_runtime_ms = None if isinstance(subblock_config, AttentionConfig): - base_runtime_ms = calc_subblock_runtime(runtime_config, None) + base_runtime_ms = _calc_llama_subblock_runtime(runtime_config, None) elif isinstance(subblock_config, FFNConfig): attn_block_config = AttentionConfig( no_op=False, num_key_value_heads=runtime_config.num_key_value_heads ).to_blockconfig() - base_runtime_ms = calc_subblock_runtime(runtime_config, attn_block_config) + base_runtime_ms = _calc_llama_subblock_runtime(runtime_config, attn_block_config) else: raise ValueError(f"Unsupported subblock type: {type(subblock_config)}") return base_runtime_ms +def _calc_moe_ffn_runtime( + runtime_config: RuntimeConfig, + ffn_config: FFNConfig, + model_config: PretrainedConfig, + descriptor: type[ModelDescriptor], + baseline: bool, +) -> float: + target_block_config = _make_moe_ffn_block_config(ffn_config) + block_config = _make_moe_ffn_baseline_block_config() if baseline else target_block_config + model = create_descriptor_benchmark_model( + model_config, + descriptor, + runtime_config.vocab_size, + runtime_config.hidden_size, + runtime_config.num_key_value_heads, + runtime_config.num_attention_heads, + runtime_config.prefill_seq_len, + runtime_config.generation_seq_len, + block_config=block_config, + repeat_block_n_times=runtime_config.repeat_block_n_times, + pattern_block_config=target_block_config, + ) + return calc_model_runtime(model, runtime_config, descriptor) + + +def calc_subblock_runtime( + runtime_config: RuntimeConfig, + subblock_config: SubblockConfig | None, + model_config: PretrainedConfig | None = None, + descriptor: type[ModelDescriptor] | None = None, +) -> float: + """Measure total runtime of a repeated subblock via vLLM latency benchmark.""" + if isinstance(subblock_config, FFNConfig) and subblock_config.is_moe: + if model_config is None or descriptor is None: + raise ValueError("MoE runtime benchmarking requires model_config and descriptor.") + return _calc_moe_ffn_runtime( + runtime_config, subblock_config, model_config, descriptor, baseline=False + ) + if _uses_moe(subblock_config): + raise ValueError(f"MoE runtime stats support FFNConfig only, got {subblock_config}.") + return _calc_llama_subblock_runtime(runtime_config, subblock_config) + + +def calc_base_runtime( + runtime_config: RuntimeConfig, + subblock_config: SubblockConfig, + model_config: PretrainedConfig | None = None, + descriptor: type[ModelDescriptor] | None = None, +) -> float: + """Calculate the base runtime of a model with no target subblocks.""" + if isinstance(subblock_config, FFNConfig) and subblock_config.is_moe: + if model_config is None or descriptor is None: + raise ValueError("MoE runtime benchmarking requires model_config and descriptor.") + return _calc_moe_ffn_runtime( + runtime_config, subblock_config, model_config, descriptor, baseline=True + ) + return _calc_llama_base_runtime(runtime_config, subblock_config) + + @cache def calc_no_block_runtime(runtime_config: RuntimeConfig) -> float: """Estimate the overhead runtime (embedding + LM head) with no decoder blocks.""" @@ -151,8 +309,8 @@ def calc_no_block_runtime(runtime_config: RuntimeConfig) -> float: block_config = _make_standard_block_config(runtime_config.num_key_value_heads) - runtime_ms_one_block = calc_subblock_runtime(runtime_config, None) # only one base block - runtime_ms_ten_blocks = calc_subblock_runtime( + runtime_ms_one_block = _calc_llama_subblock_runtime(runtime_config, None) # only one base block + runtime_ms_ten_blocks = _calc_llama_subblock_runtime( runtime_cfg_ten_blocks, block_config ) # one base block + 9 repeated blocks @@ -172,6 +330,8 @@ def calc_runtime_for_subblocks( prefill_seq_len: int, generation_seq_len: int, batch_size: int, + model_config: PretrainedConfig | None = None, + descriptor: type[ModelDescriptor] | None = None, ) -> tuple[dict[SubblockConfig, float], float]: """Benchmark each unique subblock and return per-subblock runtimes and no-block overhead.""" repeat_block_n_times = 10 @@ -197,12 +357,16 @@ def calc_runtime_for_subblocks( sorted(subblock_config_set), desc=(f"Computing runtime for {len(subblock_config_set)} subblocks\n"), ): - baseline_runtime_ms = calc_base_runtime(runtime_config, subblock_config) + baseline_runtime_ms = calc_base_runtime( + runtime_config, subblock_config, model_config, descriptor + ) if subblock_config.no_op: total_runtime_ms = 0.0 else: - subblock_total_runtime_ms = calc_subblock_runtime(runtime_config, subblock_config) + subblock_total_runtime_ms = calc_subblock_runtime( + runtime_config, subblock_config, model_config, descriptor + ) total_runtime_ms = ( subblock_total_runtime_ms - baseline_runtime_ms ) / repeat_block_n_times diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 1d04cc01add..75985e68634 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py @@ -132,6 +132,8 @@ def calculate_subblock_stats( prefill_seq_len=prefill_seq_len, generation_seq_len=generation_seq_len, batch_size=batch_size, + model_config=model_config, + descriptor=descriptor, ) sorted_subblock_config = sorted( diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index 204c2a74305..e337a59329e 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -32,6 +32,7 @@ from transformers import AutoTokenizer, LlamaForCausalLM from ..anymodel.converter import Converter +from ..anymodel.model_descriptor import ModelDescriptor from ..anymodel.models.llama import LlamaModelDescriptor from ..tools.logger import mprint from ..utils.vllm_adapter import convert_block_configs_to_per_layer_config @@ -59,10 +60,15 @@ class RuntimeConfig: gpu_memory_utilization: float = 0.5 -def save_model(model: LlamaForCausalLM, tokenizer_path: Path, output_path: Path) -> None: +def save_model( + model: LlamaForCausalLM, + tokenizer_path: Path, + output_path: Path, + descriptor: type[ModelDescriptor] = LlamaModelDescriptor, +) -> None: """Save model weights as AnyModel and copy the tokenizer to ``output_path``.""" model = model.to(dtype=torch.bfloat16) - save_model_as_anymodel(model, output_path, LlamaModelDescriptor) + save_model_as_anymodel(model, output_path, descriptor) tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) tokenizer.save_pretrained(output_path) diff --git a/tests/gpu_vllm/torch/puzzletron/test_calc_runtime_stats.py b/tests/gpu_vllm/torch/puzzletron/test_calc_runtime_stats.py index 377a2ffed19..8d2c578dc92 100644 --- a/tests/gpu_vllm/torch/puzzletron/test_calc_runtime_stats.py +++ b/tests/gpu_vllm/torch/puzzletron/test_calc_runtime_stats.py @@ -27,8 +27,35 @@ from _test_utils.torch.transformers_models import get_tiny_tokenizer from omegaconf import OmegaConf -from modelopt.torch.puzzletron.block_config import AttentionConfig, FFNConfig -from modelopt.torch.puzzletron.subblock_stats.calc_runtime_stats import calc_runtime_for_subblocks +from modelopt.torch.puzzletron.block_config import AttentionConfig, FFNConfig, MoEConfig +from modelopt.torch.puzzletron.subblock_stats.calc_runtime_stats import ( + _hybrid_override_pattern_for_block_configs, + _make_moe_ffn_baseline_block_config, + _make_moe_ffn_block_config, + _uses_moe, + calc_runtime_for_subblocks, +) + + +def test_moe_ffn_runtime_path_preserves_explicit_moe_block(): + """MoE expert-removal candidates should benchmark as real MoE layers.""" + pruned_ffn = FFNConfig( + moe=MoEConfig( + num_local_experts=16, + num_experts_per_tok=4, + expert_intermediate_dim=128, + shared_expert_intermediate_dim=256, + ), + ) + + target_block = _make_moe_ffn_block_config(pruned_ffn) + baseline_block = _make_moe_ffn_baseline_block_config() + + assert _uses_moe(pruned_ffn) + assert target_block.attention.no_op + assert target_block.ffn.moe.num_local_experts == 16 + assert baseline_block.attention.no_op and baseline_block.ffn.no_op + assert _hybrid_override_pattern_for_block_configs([target_block]) == "E" @pytest.mark.skip(reason="AnyModel is not supported in vLLM yet")