From a724afd0ee38c7a62f8206bf2cc11dcaa4fe8685 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Mon, 15 Jun 2026 12:09:27 -0700 Subject: [PATCH 1/9] 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/9] 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/9] 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/9] 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/9] 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/9] 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 5d5e44bd54adcc980f0de87aebc7caeea3712642 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Thu, 9 Jul 2026 03:27:50 -0700 Subject: [PATCH 7/9] added parameters to control mutli-gpu benchmarking Signed-off-by: Grzegorz Karch --- .../llama-3_1-8B_pruneffn_runtime.yaml | 5 +++++ .../subblock_stats/calc_runtime_stats.py | 5 +++++ .../puzzletron/subblock_stats/runtime_utils.py | 5 +++++ .../puzzletron/subblock_stats/runtime_vllm.py | 15 +++++++++++---- 4 files changed, 26 insertions(+), 4 deletions(-) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml index 2ca0f2c16cf..06fbb53e118 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml @@ -16,6 +16,11 @@ calc_subblock_stats: enabled: true num_warmup_iters: 2 num_iters: 10 + data_parallel_size: 1 + tensor_parallel_size: 1 + prefill_context_parallel_size: 1 + decode_context_parallel_size: 1 + enable_expert_parallel: false # FFN intermediate sizes to search over (heterogeneous architecture) pruning: diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py index c7d926de18f..c810ae64fed 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py @@ -186,6 +186,11 @@ def calc_runtime_for_subblocks( prefill_seq_len, generation_seq_len, batch_size, + runtime_stats_config.get("data_parallel_size", 1), + runtime_stats_config.get("tensor_parallel_size", 1), + runtime_stats_config.get("prefill_context_parallel_size", 1), + runtime_stats_config.get("decode_context_parallel_size", 1), + runtime_stats_config.get("enable_expert_parallel", False), runtime_stats_config.get("num_iters", 30), runtime_stats_config.get("num_warmup_iters", 10), runtime_stats_config.get("gpu_memory_utilization", 0.5), diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py index 735c3f8f72b..949c980146a 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -49,6 +49,11 @@ class RuntimeConfig: prefill_seq_len: int generation_seq_len: int batch_size: int + data_parallel_size: int + tensor_parallel_size: int + prefill_context_parallel_size: int + decode_context_parallel_size: int + enable_expert_parallel: bool num_iters: int num_warmup_iters: int # Fraction of total GPU memory vLLM may use. Kept well below the default diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py index 386f3615d49..364fcd314dc 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py @@ -79,10 +79,17 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) str(runtime_config.num_iters), "--max-num-seqs", "1", + "--data-parallel-size", + str(runtime_config.data_parallel_size), + "--prefill-context-parallel-size", + str(runtime_config.prefill_context_parallel_size), + "--decode-context-parallel-size", + str(runtime_config.decode_context_parallel_size), + "--enable-expert-parallel" + if runtime_config.enable_expert_parallel + else "--no-enable-expert-parallel", "--tensor-parallel-size", - "1", - "--pipeline-parallel-size", - "1", + str(runtime_config.tensor_parallel_size), "--distributed-executor-backend", "external_launcher", # Cap GPU memory so vLLM's startup free-memory check passes while the @@ -101,7 +108,7 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) check=True, capture_output=True, text=True, - timeout=1800, # 30 minutes + timeout=3600, # 1 hour ) # nosec B603 except subprocess.TimeoutExpired as exc: raise TimeoutError("vLLM latency benchmark timed out") from exc From c6d55932c909024d73078c8fa1f3e9914296d5e4 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Tue, 14 Jul 2026 05:42:44 -0700 Subject: [PATCH 8/9] added extra_args to pass to vlmm Signed-off-by: Grzegorz Karch --- .../llama-3_1-8B_pruneffn_runtime.yaml | 1 + .../subblock_stats/calc_runtime_stats.py | 1 + .../subblock_stats/runtime_utils.py | 1 + .../puzzletron/subblock_stats/runtime_vllm.py | 24 ++++++++++++++++++- 4 files changed, 26 insertions(+), 1 deletion(-) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml index 06fbb53e118..e788443defd 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_runtime/llama-3_1-8B_pruneffn_runtime.yaml @@ -21,6 +21,7 @@ calc_subblock_stats: prefill_context_parallel_size: 1 decode_context_parallel_size: 1 enable_expert_parallel: false + extra_args: "" # FFN intermediate sizes to search over (heterogeneous architecture) pruning: diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py index c810ae64fed..669b25e675d 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py +++ b/modelopt/torch/puzzletron/subblock_stats/calc_runtime_stats.py @@ -194,6 +194,7 @@ def calc_runtime_for_subblocks( 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_stats_config.get("extra_args", ""), ) 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 949c980146a..814d08966de 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_utils.py @@ -61,6 +61,7 @@ class RuntimeConfig: # GPU during benchmarking; requesting too much fails vLLM's startup # free-memory check. gpu_memory_utilization: float = 0.5 + extra_args: str = "" def save_model(model: LlamaForCausalLM, tokenizer_path: Path, output_path: Path) -> None: diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py index 364fcd314dc..8d5d18453d7 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py @@ -27,6 +27,7 @@ """ import json +import os import subprocess # nosec B404 from pathlib import Path from types import SimpleNamespace @@ -91,7 +92,7 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) "--tensor-parallel-size", str(runtime_config.tensor_parallel_size), "--distributed-executor-backend", - "external_launcher", + "mp", #"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", @@ -100,6 +101,26 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) "--optimization-level", "0", ] + if runtime_config.extra_args: + cmd.extend(runtime_config.extra_args.split(" ")) + + # Strip torchrun-injected distributed vars so vLLM's mp workers don't + # inherit an outer process group's rendezvous state, which causes a + # deadlock when vLLM tries to create its own Gloo CPU group. + _TORCHRUN_VARS = { + "RANK", + "LOCAL_RANK", + "WORLD_SIZE", + "LOCAL_WORLD_SIZE", + "MASTER_ADDR", + "MASTER_PORT", + "TORCHELASTIC_RESTART_COUNT", + "TORCHELASTIC_MAX_RESTARTS", + "TORCHELASTIC_RUN_ID", + "TORCHELASTIC_USE_AGENT_STORE", + "TORCHELASTIC_ERROR_FILE", + } + env = {k: v for k, v in os.environ.items() if k not in _TORCHRUN_VARS} # cmd is a fixed list of strings (no shell, no untrusted input). try: @@ -109,6 +130,7 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) capture_output=True, text=True, timeout=3600, # 1 hour + env=env, ) # nosec B603 except subprocess.TimeoutExpired as exc: raise TimeoutError("vLLM latency benchmark timed out") from exc From 509cd231b8e641bcb6426ac7a99cb86c5819e948 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Tue, 14 Jul 2026 05:46:00 -0700 Subject: [PATCH 9/9] formating Signed-off-by: Grzegorz Karch --- modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py index 8d5d18453d7..bd2493d6115 100644 --- a/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py +++ b/modelopt/torch/puzzletron/subblock_stats/runtime_vllm.py @@ -92,7 +92,7 @@ def run_vllm_latency_benchmark(model_path: Path, runtime_config: RuntimeConfig) "--tensor-parallel-size", str(runtime_config.tensor_parallel_size), "--distributed-executor-backend", - "mp", #"external_launcher", + "mp", # "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",