ple_pruning#

Global per-layer-embedding (PLE) channel ranking and tensor surgery.

Classes

PLEPruningSpec

Own every tensor coupled to one global PLE channel dimension.
class PLEPruningSpec#

Bases: object

Own every tensor coupled to one global PLE channel dimension.

Gemma-style PLE uses one shared RMSNorm across all layer chunks. Therefore every layer must use the same permutation and retained width. Per-layer contribution scores are summed before sorting to preserve this invariant.

__init__(language_prefix, layer_template, num_layers, width, layer_gate_name='per_layer_input_gate', layer_projection_name='per_layer_projection', model_embedding_name='embed_tokens_per_layer', model_projection_name='per_layer_model_projection', model_norm_name='per_layer_projection_norm')#
Parameters:

  • language_prefix (str)

  • layer_template (str)

  • num_layers (int)

  • width (int)

  • layer_gate_name (str)

  • layer_projection_name (str)

  • model_embedding_name (str)

  • model_projection_name (str)

  • model_norm_name (str)

Return type:

None

language_prefix: str#
layer_gate_name: str = 'per_layer_input_gate'#
layer_prefix(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

layer_projection_name: str = 'per_layer_projection'#
layer_score_key(layer_idx)#
Parameters:

layer_idx (int)

Return type:

str

layer_template: str#
model_embedding_name: str = 'embed_tokens_per_layer'#
model_norm_name: str = 'per_layer_projection_norm'#
model_projection_name: str = 'per_layer_model_projection'#
num_layers: int#
order_from_score_logs(score_logs)#
Parameters:

score_logs (dict[str, dict])

Return type:

Tensor

permute_state_dict(state_dict, order)#
Parameters:

  • state_dict (dict[str, Tensor])

  • order (Tensor)

Return type:

tuple[dict[str, Tensor], set[str]]

slice_state_dict(state_dict, target_width)#
Parameters:

  • state_dict (dict[str, Tensor])

  • target_width (int)

Return type:

dict[str, Tensor]

width: int#