embedding_pruning#

Descriptor-owned residual hidden-width ranking, sorting, and slicing.

Classes

EmbeddingPruningSpec

Complete model-family contract for the language residual width.

PackedMinitronImportance

Exact Minitron hidden metric with original packed-sample boundaries.

TensorAxisRule

class EmbeddingPruningSpec#

Bases: object

Complete model-family contract for the language residual width.

__init__(hidden_size, legal_widths, alignment, tensor_rules, exempt_patterns=(), tie_groups=(), config_paths=(('hidden_size',),), residual_norm_patterns=(), permutation_group_size=1)#
Parameters:

  • hidden_size (int)

  • legal_widths (tuple[int, ...])

  • alignment (int)

  • tensor_rules (tuple[TensorAxisRule, ...])

  • exempt_patterns (tuple[str, ...])

  • tie_groups (tuple[tuple[str, ...], ...])

  • config_paths (tuple[tuple[str, ...], ...])

  • residual_norm_patterns (tuple[str, ...])

  • permutation_group_size (int)

Return type:

None

alignment: int#
audit_state_dict(state_dict)#
Parameters:

state_dict (Mapping[str, Tensor])

Return type:

dict[str, Any]

config_paths: tuple[tuple[str, ...], ...] = (('hidden_size',),)#
exempt_patterns: tuple[str, ...] = ()#
hidden_size: int#
legal_widths: tuple[int, ...]#
order_from_scores(scores)#

Return a legal residual permutation for the descriptor’s storage granularity.

Block-quantized inputs cannot move individual channels across quantization blocks without dequantizing and requantizing. Such descriptors rank whole blocks by summed channel importance and preserve channel order within a block.

Parameters:

scores (Tensor)

Return type:

Tensor

parameter_count(state_dict)#
Parameters:

state_dict (Mapping[str, Tensor])

Return type:

int

permutation_group_size: int = 1#
permute_state_dict(state_dict, order)#
Parameters:

  • state_dict (Mapping[str, Tensor])

  • order (Tensor)

Return type:

dict[str, Tensor]

residual_norm_patterns: tuple[str, ...] = ()#
rule_for(key)#

Return the unique descriptor rule for a canonical tensor key.

Parameters:

key (str)

Return type:

TensorAxisRule | None

slice_state_dict(state_dict, width, *, tp_size=1)#
Parameters:

  • state_dict (Mapping[str, Tensor])

  • width (int)

  • tp_size (int)

Return type:

dict[str, Tensor]

sliced_shape(key, shape, width)#

Return a tensor’s exact shape after residual-width slicing.

Parameters:

  • key (str)

  • shape (Sequence[int])

  • width (int)

Return type:

tuple[int, …]

tensor_rules: tuple[TensorAxisRule, ...]#
tie_groups: tuple[tuple[str, ...], ...] = ()#
update_config(config, width)#
Parameters:

  • config (Mapping[str, Any])

  • width (int)

Return type:

dict[str, Any]

update_config_object(config, width)#

Return a deep-copied config with every descriptor-owned width updated.

AutoModel/HF configuration trees mix PretrainedConfig objects and dictionaries. Physical width realization must update both the nested language config and any top-level mirror without round-tripping through to_dict(), which would lose the concrete config classes needed by model construction.

Parameters:

  • config (Any)

  • width (int)

Return type:

Any

validate_width(width, *, tp_size=1)#
Parameters:

  • width (int)

  • tp_size (int)

Return type:

int

class PackedMinitronImportance#

Bases: object

Exact Minitron hidden metric with original packed-sample boundaries.

__init__(hidden_size)#
Parameters:

hidden_size (int)

property sample_count: int#
scores()#
Return type:

Tensor

property site_names: tuple[str, ...]#
squared_sums()#
Return type:

dict[str, Tensor]

update(site_name, activations, *, cu_seqlens=None, token_mask=None)#
Parameters:

  • site_name (str)

  • activations (Tensor)

  • cu_seqlens (Tensor | None)

  • token_mask (Tensor | None)

Return type:

None

class TensorAxisRule#

Bases: object

__init__(pattern, axes, description, chunked_axes=(), grouped_axes=())#
Parameters:

  • pattern (str)

  • axes (tuple[int, ...])

  • description (str)

  • chunked_axes (tuple[tuple[int, int], ...])

  • grouped_axes (tuple[tuple[int, int], ...])

Return type:

None

axes: tuple[int, ...]#
chunked_axes: tuple[tuple[int, int], ...] = ()#
description: str#
grouped_axes: tuple[tuple[int, int], ...] = ()#
matches(key)#
Parameters:

key (str)

Return type:

bool

pattern: str#