pruning_ckpts

Utilities for creating pruned model checkpoints.

This module provides functions to generate pruned checkpoints by modifying model architectures (FFN intermediate sizes, attention head groups, hidden dimensions) and initializing child pruned models from parent checkpoints.

Functions

launch_ffn_intermediates_prune_ckpt

launch_attn_groups_prune_ckpt

launch_hidden_dim_prune_ckpt

Launch hidden dimension pruning using channel importance ranking.

launch_experts_prune_ckpt

launch_moe_ffn_intermediates_prune_ckpt

launch_prune_ckpt

launch_attn_groups_prune_ckpt(cfg, max_save_workers=None, max_layer_workers=None)
Parameters:
  • cfg (DictConfig)

  • max_save_workers (int | None)

  • max_layer_workers (int | None)

launch_experts_prune_ckpt(cfg, max_save_workers=None, max_layer_workers=None, symlink_suffix=None)
Parameters:
  • cfg (DictConfig)

  • max_save_workers (int | None)

  • max_layer_workers (int | None)

  • symlink_suffix (str | None)

launch_ffn_intermediates_prune_ckpt(cfg, max_save_workers=None, max_layer_workers=None)
Parameters:
  • cfg (DictConfig)

  • max_save_workers (int | None)

  • max_layer_workers (int | None)

launch_hidden_dim_prune_ckpt(cfg)

Launch hidden dimension pruning using channel importance ranking.

Parameters:

cfg (DictConfig)

launch_moe_ffn_intermediates_prune_ckpt(cfg, max_save_workers=None, max_layer_workers=None)
Parameters:
  • cfg (DictConfig)

  • max_save_workers (int | None)

  • max_layer_workers (int | None)

launch_prune_ckpt(cfg)
Parameters:

cfg (DictConfig)