Puzzletron tutorial fixes for runtime optimization

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 <model_dir>
 ```
 
+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

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

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"

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"

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"

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}

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: validation
+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}

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

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 = {}

modelopt/torch/puzzletron/subblock_stats/runtime_utils.py

@@ -23,14 +23,18 @@
 """
 
 import json
+import shutil
 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 +52,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 +89,40 @@ 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(config_dir: Path):
+    """Convert a model to vLLM AnyModel format."""
+    # Load the model config.json, update "architectures" to ["AnyModel"], and write back to disk.
+    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(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"  # TODO: extend support to other models
+
+    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(config_path, "w") as f:
+        json.dump(vars(config), f, indent=2)
+
+
+if __name__ == "__main__":
+    import fire
+
+    fire.Fire()

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",