From a724afd0ee38c7a62f8206bf2cc11dcaa4fe8685 Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Mon, 15 Jun 2026 12:09:27 -0700 Subject: [PATCH 1/7] 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/7] 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/7] 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/7] 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/7] 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/7] 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/7] 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()