sorted_teacher

Build a “sorted teacher”: the teacher with each layer’s FFN channels and attention heads reordered by importance (from activation scoring).

Because the reorder is a consistent permutation of a contracted dim (see attention_ffn_surgery), the sorted teacher computes the same outputs as the teacher — it is the same model, reindexed so that “most important first”. Pruning any variant then becomes a prefix slice (FFN [:K]; attention = first kv groups x first m query heads), which is what lets the block library, replace-1-block scoring, bypass, and final realization all operate on this single artifact with no realized intermediate checkpoints.

This module is backend-agnostic (a weight-only transform): the scores come from the unified AutoModel activation pass, and the per-layer key layout comes from the descriptor.

Classes

LayerLayout

Per-layer weight keys + head geometry the sorter needs (descriptor-derived).

Functions

iter_safetensor_weight_files	Return model safetensor files relative to a standard HF checkpoint directory.
sort_state_dict	Return (sorted_state_dict, permutations) — the keystone, pure transform.
build_layer_layouts	Build per-layer layouts from the (cast) block configs + a layer-prefix template.
build_sorted_teacher	Write the sorted teacher to output_dir and return it.

class LayerLayout

Bases: object

Per-layer weight keys + head geometry the sorter needs (descriptor-derived).

__init__(layer_idx, head_dim, gate_key=None, up_key=None, down_key=None, ffn_intermediate=None, q_key=None, k_key=None, v_key=None, o_key=None, num_q_heads=None, num_kv_heads=None, q_gate_row_group=None, q_head_aux_keys=(), mla_prefix=None, mla_q_a_key=None, mla_q_norm_key=None, mla_q_b_key=None, mla_kv_a_key=None, mla_kv_norm_key=None, mla_kv_b_key=None, mla_o_key=None, mla_num_heads=None, mla_q_lora_rank=None, mla_kv_lora_rank=None, moe_prefix=None, moe_gate_prefix=None, moe_gate_key=None, moe_gate_bias_key=None, moe_router_aux_keys=(), moe_experts_prefix=None, moe_expert_up_keys=None, moe_expert_down_keys=None, moe_fused_expert_keys=(), moe_fused_gate_up_keys=(), moe_fused_down_keys=(), moe_expert_intermediate_group_size=1, moe_expert_order_mode='physical', moe_fused_gate_layout='concatenated', moe_shared_gate_key=None, moe_shared_up_key=None, moe_shared_down_key=None, moe_fc1_latent_key=None, moe_fc2_latent_key=None, moe_num_experts=None, moe_expert_intermediate=None, moe_shared_intermediate=None, moe_latent_dim=None, mamba_prefix=None, mamba_in_key=None, mamba_out_key=None, mamba_conv_key=None, mamba_conv_bias_key=None, mamba_a_key=None, mamba_d_key=None, mamba_dt_bias_key=None, mamba_norm_key=None, mamba_num_heads=None, mamba_head_dim=None, mamba_num_groups=None, mamba_state_dim=None, gated_delta_net=False)

Parameters:

• layer_idx (int)

• head_dim (int)

• gate_key (str | None)

• up_key (str | None)

• down_key (str | None)

• ffn_intermediate (int | None)

• q_key (str | None)

• k_key (str | None)

• v_key (str | None)

• o_key (str | None)

• num_q_heads (int | None)

• num_kv_heads (int | None)

• q_gate_row_group (int | None)

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

• mla_prefix (str | None)

• mla_q_a_key (str | None)

• mla_q_norm_key (str | None)

• mla_q_b_key (str | None)

• mla_kv_a_key (str | None)

• mla_kv_norm_key (str | None)

• mla_kv_b_key (str | None)

• mla_o_key (str | None)

• mla_num_heads (int | None)

• mla_q_lora_rank (int | None)

• mla_kv_lora_rank (int | None)

• moe_prefix (str | None)

• moe_gate_prefix (str | None)

• moe_gate_key (str | None)

• moe_gate_bias_key (str | None)

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

• moe_experts_prefix (str | None)

• moe_expert_up_keys (list[str] | None)

• moe_expert_down_keys (list[str] | None)

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

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

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

• moe_expert_intermediate_group_size (int)

• moe_expert_order_mode (str)

• moe_fused_gate_layout (str)

• moe_shared_gate_key (str | None)

• moe_shared_up_key (str | None)

• moe_shared_down_key (str | None)

• moe_fc1_latent_key (str | None)

• moe_fc2_latent_key (str | None)

• moe_num_experts (int | None)

• moe_expert_intermediate (int | None)

• moe_shared_intermediate (int | None)

• moe_latent_dim (int | None)

• mamba_prefix (str | None)

• mamba_in_key (str | None)

• mamba_out_key (str | None)

• mamba_conv_key (str | None)

• mamba_conv_bias_key (str | None)

• mamba_a_key (str | None)

• mamba_d_key (str | None)

• mamba_dt_bias_key (str | None)

• mamba_norm_key (str | None)

• mamba_num_heads (int | None)

• mamba_head_dim (int | None)

• mamba_num_groups (int | None)

• mamba_state_dim (int | None)

• gated_delta_net (bool)

Return type:

None

down_key: str | None = None

ffn_intermediate: int | None = None

gate_key: str | None = None

gated_delta_net: bool = False

head_dim: int

k_key: str | None = None

layer_idx: int

mamba_a_key: str | None = None

mamba_conv_bias_key: str | None = None

mamba_conv_key: str | None = None

mamba_d_key: str | None = None

mamba_dt_bias_key: str | None = None

mamba_head_dim: int | None = None

mamba_in_key: str | None = None

mamba_norm_key: str | None = None

mamba_num_groups: int | None = None

mamba_num_heads: int | None = None

mamba_out_key: str | None = None

mamba_prefix: str | None = None

mamba_state_dim: int | None = None

mla_kv_a_key: str | None = None

mla_kv_b_key: str | None = None

mla_kv_lora_rank: int | None = None

mla_kv_norm_key: str | None = None

mla_num_heads: int | None = None

mla_o_key: str | None = None

mla_prefix: str | None = None

mla_q_a_key: str | None = None

mla_q_b_key: str | None = None

mla_q_lora_rank: int | None = None

mla_q_norm_key: str | None = None

moe_expert_down_keys: list[str] | None = None

moe_expert_intermediate: int | None = None

moe_expert_intermediate_group_size: int = 1

moe_expert_order_mode: str = 'physical'

moe_expert_up_keys: list[str] | None = None

moe_experts_prefix: str | None = None

moe_fc1_latent_key: str | None = None

moe_fc2_latent_key: str | None = None

moe_fused_down_keys: tuple[str, ...] = ()

moe_fused_expert_keys: tuple[str, ...] = ()

moe_fused_gate_layout: str = 'concatenated'

moe_fused_gate_up_keys: tuple[str, ...] = ()

moe_gate_bias_key: str | None = None

moe_gate_key: str | None = None

moe_gate_prefix: str | None = None

moe_latent_dim: int | None = None

moe_num_experts: int | None = None

moe_prefix: str | None = None

moe_router_aux_keys: tuple[str, ...] = ()

moe_shared_down_key: str | None = None

moe_shared_gate_key: str | None = None

moe_shared_intermediate: int | None = None

moe_shared_up_key: str | None = None

num_kv_heads: int | None = None

num_q_heads: int | None = None

o_key: str | None = None

q_gate_row_group: int | None = None

q_head_aux_keys: tuple[str, ...] = ()

q_key: str | None = None

up_key: str | None = None

v_key: str | None = None

build_layer_layouts(block_configs, *, layer_prefix_tmpl, num_attention_heads, head_dim, ffn_subnames=('gate_proj', 'up_proj', 'down_proj'), attn_subnames=('q_proj', 'k_proj', 'v_proj', 'o_proj'), mlp_module='mlp', attn_module='self_attn', moe_module='mlp', moe_router_subname='gate', moe_router_aux_subnames=(), moe_fused_expert_subnames=(), moe_fused_gate_up_subnames=(), moe_fused_down_subnames=(), moe_expert_intermediate_group_size=1, moe_expert_order_mode='physical', moe_fused_gate_layout='concatenated', moe_shared_expert_subname='shared_experts', moe_shared_gate_subname=None, moe_shared_up_subname='up_proj', moe_shared_down_subname='down_proj', mamba_module='self_attn', q_gate_row_group=None, attention_q_head_subnames=(), gated_delta_net=False)

Build per-layer layouts from the (cast) block configs + a layer-prefix template.

layer_prefix_tmpl is e.g. "model.language_model.layers.{i}". FFN/attention keys use the standard HF sub-names (overridable). A layer contributes FFN keys only if its FFN is prunable (not no-op) and attention keys only if its attention is a real GQA block (num_kv_heads set, i.e. not mamba / no-op).

Parameters:

• layer_prefix_tmpl (str)

• num_attention_heads (int)

• head_dim (int)

• ffn_subnames (tuple[str, str, str])

• attn_subnames (tuple[str, str, str, str])

• mlp_module (str)

• attn_module (str)

• moe_module (str)

• moe_router_subname (str)

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

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

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

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

• moe_expert_intermediate_group_size (int)

• moe_expert_order_mode (str)

• moe_fused_gate_layout (str)

• moe_shared_expert_subname (str)

• moe_shared_gate_subname (str | None)

• moe_shared_up_subname (str)

• moe_shared_down_subname (str)

• mamba_module (str)

• q_gate_row_group (int | None)

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

• gated_delta_net (bool)

Return type:

list[LayerLayout]

build_sorted_teacher(teacher_dir, activations_log_dir, output_dir, descriptor, *, ffn_subnames=('gate_proj', 'up_proj', 'down_proj'), attn_subnames=('q_proj', 'k_proj', 'v_proj', 'o_proj'), mlp_module='mlp', attn_module='self_attn', deferred_axes=(), mamba_state_score_key='ssm_channel_contrib', embedding_widths=None)

Write the sorted teacher to output_dir and return it.

Copies every non-weight file of the AnyModel checkpoint, then rewrites the standard HuggingFace safetensors shards with FFN/attention tensors permuted by importance (keys/shapes/index unchanged, so the index stays valid and the model is functionally identical). Also writes sorted_permutations.json.

Parameters:

• teacher_dir (str | Path)

• activations_log_dir (str | Path)

• output_dir (str | Path)

• ffn_subnames (tuple[str, str, str])

• attn_subnames (tuple[str, str, str, str])

• mlp_module (str)

• attn_module (str)

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

• mamba_state_score_key (str)

• embedding_widths (Sequence[int] | None)

Return type:

Path

iter_safetensor_weight_files(checkpoint_dir)

Return model safetensor files relative to a standard HF checkpoint directory.

Puzzletron canonical checkpoints are standard HuggingFace artifacts.

Parameters:

checkpoint_dir (str | Path)

Return type:

tuple[Path, …]

sort_state_dict(state_dict, layouts, ffn_scores, attention_logs, score_logs=None, original_tensor_loader=None, latent_transform_cache=None, deferred_axes=frozenset({}), mamba_state_score_key='ssm_channel_contrib')

Return (sorted_state_dict, permutations) — the keystone, pure transform.

ffn_scores[i] is the per-intermediate-channel importance [I]. attention_logs[i] carries explicit grouped-attention scores: kv_group_scores [num_kv] and query_head_scores [num_kv, heads_per_group]. Only keys present in state_dict are touched, so this is safe for heterogeneous (mamba / no-op) layers. permutations records the applied orders.

Parameters:

• state_dict (dict[str, Tensor])

• layouts (list[LayerLayout])

• ffn_scores (dict[int, Tensor])

• attention_logs (dict[int, dict])

• score_logs (dict[str, dict] | None)

• original_tensor_loader (Callable[[str], Tensor | None] | None)

• latent_transform_cache (dict[int, LatentMoETransform] | None)

• deferred_axes (frozenset[str])

• mamba_state_score_key (str)

Return type:

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