From d8c2f8c85407f618cfa7edc7d48a1f16106b3f1a Mon Sep 17 00:00:00 2001 From: Grzegorz Karch Date: Tue, 28 Apr 2026 04:07:46 -0700 Subject: [PATCH] added required keys Signed-off-by: Grzegorz Karch --- .../llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml index 01886607e4d..53d7e4bd9c6 100644 --- a/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml +++ b/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/pruning/attn_pruning.yaml @@ -1,8 +1,15 @@ 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