materialize

Materialize a pruned variant from the sorted teacher — slice/merge, no realized checkpoints.

This is the realize-side counterpart of sorted_teacher.sort_state_dict() and the static twin of the dynamic prune (dynamic_block_prune): given the sorted teacher’s state dict and a per-layer target (ffn K / attn (q, kv)), it produces the physically smaller weights for the chosen variant by prefix-slicing (removal) or merging (q-preserving KV reduction). It backs the block library (a variant = a slice/merge spec into the sorted teacher, not realized weight_paths) and the bypass / final-realize steps.

Because the sorted teacher orders channels/heads by importance, slicing [:K] (and the blocked attention keep-set) yields exactly the importance-pruned weights. Reuses the model-agnostic primitives in attention_ffn_surgery; the descriptor supplies the per-layer keys/geometry via sorted_teacher.build_layer_layouts().

Classes

BlockTarget

Target dims for one layer's replaced subblock (None = unchanged for that axis).

Functions

`materialize_solution_state_dict`	Return a new state dict with each targeted layer's block sliced/merged to its variant.
`block_targets_from_replacements`	Convert replacement-library `layer_replacements` -> `{layer_idx: BlockTarget}`.
`materialize_model_from_sorted`	Build a pruned model by materializing `layer_replacements` from the sorted teacher.
`materialize_checkpoint_from_sorted`	Stream one sorted-teacher shard at a time into a realized HF checkpoint.
`materialize_hidden_width_checkpoint`	Physically realize only the descriptor-owned residual width.

class BlockTarget

Bases: object

Target dims for one layer’s replaced subblock (None = unchanged for that axis).

__init__(target_intermediate=None, target_num_q=None, target_num_kv=None, target_mla_heads=None, target_q_lora_rank=None, target_kv_lora_rank=None, target_num_experts=None, expert_keep_indices=None, target_expert_intermediate=None, target_shared_expert_intermediate=None, target_latent_dim=None, target_mamba_heads=None, target_mamba_groups=None, target_mamba_head_dim=None, target_mamba_state_dim=None, remove_ffn=False, remove_attention=False, remove_mamba=False, remove_moe=False)

Parameters:

target_intermediate (int | None)
target_num_q (int | None)
target_num_kv (int | None)
target_mla_heads (int | None)
target_q_lora_rank (int | None)
target_kv_lora_rank (int | None)
target_num_experts (int | None)
expert_keep_indices (tuple[int, ...] | None)
target_expert_intermediate (int | None)
target_shared_expert_intermediate (int | None)
target_latent_dim (int | None)
target_mamba_heads (int | None)
target_mamba_groups (int | None)
target_mamba_head_dim (int | None)
target_mamba_state_dim (int | None)
remove_ffn (bool)
remove_attention (bool)
remove_mamba (bool)
remove_moe (bool)

Return type:

None

expert_keep_indices: tuple[int, ...] | None = None

remove_attention: bool = False

remove_ffn: bool = False

remove_mamba: bool = False

remove_moe: bool = False

target_expert_intermediate: int | None = None

target_intermediate: int | None = None

target_kv_lora_rank: int | None = None

target_latent_dim: int | None = None

target_mamba_groups: int | None = None

target_mamba_head_dim: int | None = None

target_mamba_heads: int | None = None

target_mamba_state_dim: int | None = None

target_mla_heads: int | None = None

target_num_experts: int | None = None

target_num_kv: int | None = None

target_num_q: int | None = None

target_q_lora_rank: int | None = None

target_shared_expert_intermediate: int | None = None

block_targets_from_replacements(layer_replacements, teacher_block_configs, num_attention_heads)

Convert replacement-library layer_replacements -> {layer_idx: BlockTarget}.

Each replacement pairs parent_layer_indices with child_block_configs (parsed BlockConfig objects). The child’s FFN intermediate_size / attention num_query_heads + num_kv_heads become the per-layer target dims that materialize_solution_state_dict() slices. If a child only reduces num_kv_heads, materialization also reduces num_query_heads by removing the corresponding sorted query-head groups.

Parameters:: num_attention_heads (int)
Return type:: dict[int, BlockTarget]

materialize_checkpoint_from_sorted(sorted_dir, layer_replacements, descriptor, child_model_config, output_dir, *, overwrite=False, solution_identity=None)

Stream one sorted-teacher shard at a time into a realized HF checkpoint.

Parameters:

overwrite (bool)
solution_identity (str | None)

Return type:

Path

materialize_hidden_width_checkpoint(sorted_dir, descriptor, hidden_width, output_dir, *, alignment=1, overwrite=False)

Physically realize only the descriptor-owned residual width.

This creates the width-specific parent used by nested replacement scoring: block configurations remain identical to the sorted teacher while every language residual tensor, projector output, LM-head input, and MTP residual axis is sliced together.

Parameters:

hidden_width (int)
alignment (int)
overwrite (bool)

Return type:

Path

materialize_model_from_sorted(sorted_dir, layer_replacements, descriptor, child_model_config)

Build a pruned model by materializing layer_replacements from the sorted teacher.

The realize/library counterpart of the dynamic prune: load the sorted teacher’s weights, slice/ merge the replaced blocks to their target dims, then load the result into a model built from child_model_config (the solution’s heterogeneous block configs). No realized intermediate checkpoint. For large models this materializes densely — the sharded path writes the result via the normal save and reloads with load_and_shard_model (an in-container follow-up).

materialize_solution_state_dict(state_dict, layouts, targets)

Return a new state dict with each targeted layer’s block sliced/merged to its variant.

state_dict is the (sorted) teacher’s; targets maps layer_idx -> BlockTarget. Untargeted layers are copied unchanged. The result has smaller tensors for materialized blocks, so it loads into a model configured with the corresponding (heterogeneous) block configs.

Parameters:: targets (dict[int, BlockTarget])