added required keys

Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>

Author: grzegorz-k-karch

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