base

Classes

ModelDescriptor

class ModelDescriptor

Bases: ABC

classmethod adapt_loaded_state_dict_for_model(state_dict, *, model, config, checkpoint_to_model_key=None)

Rewrite loaded checkpoint keys to the live model’s state_dict names.

Parameters:

• state_dict (dict[str, Any])

• model (Module)

• config (Any)

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

Return type:

dict[str, Any]

classmethod adapt_materialized_state_dict_for_model(state_dict, *, model, config)

Rewrite realized checkpoint keys before loading into a freshly built model.

The sorted teacher keeps Puzzletron’s canonical AnyModel checkpoint names. Some native HF/AutoModel classes expose the same weights under different module prefixes during initialization. Descriptors can override this hook to bridge that boundary without teaching the generic materializer about family-specific names.

Parameters:

• state_dict (dict[str, Any])

• model (Module)

• config (Any)

Return type:

dict[str, Any]

classmethod adapt_module_name_for_model(module_name, model)

Rewrite a descriptor module FQN for an instantiated runtime model if needed.

Descriptor names are Puzzletron’s canonical checkpoint names. Some remote-code classes expose the same modules under different initialization prefixes. Keep that compatibility at the descriptor boundary so sharding and native runtime views can address live modules without changing canonical checkpoint keys.

Parameters:

• module_name (str)

• model (Module)

Return type:

str

classmethod anymodel_arch_info()

Return vLLM AnyModel architecture metadata for this descriptor.

The vLLM fork can either use its built-in registry keyed by base_architecture or a config-local anymodel_arch_info contract. Runtime benchmark checkpoints are synthetic, so emit the minimal config-local contract by default from the HF decoder layer class. Hybrid/custom families should override this with their full contract.

Return type:

dict[str, Any]

static attn_no_op_post_init(decoder_layer)

Post-init callback to alter a decoder layer so that Attention subblock performs as no-op.

It is recommended to use the utils modules from no_op.py to replace layers to dummy counterparts.

Example for replacing a layernorm layer with identity:

>>> decoder_layer.post_attention_layernorm = Same()

Example for replacing an attention layer with zeroes:

>>> decoder_layer.self_attn = MatchingZeros()

In case the attention layer returns multiple outputs i.e hidden_states, _ = self.self_attn(), use the util method return_tuple_of_size to return trailing None values:

>>> decoder_layer.self_attn = return_tuple_of_size(MatchingZeros, size=2)()

Parameters:

decoder_layer (Module)

classmethod attn_no_op_supported()

Check whether attn_no_op_post_init is overridden for attention no-op support.

classmethod automodel_model_kwargs(config, *, distributed=None)

Model-construction settings required by the native AutoModel family.

Parameters:

distributed (dict[str, Any] | None)

Return type:

dict[str, Any]

classmethod automodel_tp_linear_backend(config)

Linear backend compatible with the native model’s TP plan.

AutoModel’s standard tensor-parallel plans use PyTorch DTensor RowwiseParallel/ColwiseParallel styles, which operate on nn.Linear rather than Transformer Engine Linear modules. A family with a custom TE-aware TP implementation may override this hook.

Return type:

str | None

abstract static block_config_to_layer_overrides(block_config)

Map between BlockConfig and layer config overrides.

These overrides are consumed by a specific decoder layer and by the whole model. Usage can be seen in deci_x_patcher under the method _patched_decoder_layer_init.

Example implementation to override the FFN intermediate size of a block:

>>> def block_config_to_layer_overrides(block_config: BlockConfig) -> Dict[str, Any]:
>>>     ffn = block_config.require_subblock("ffn")
>>>     return {"intermediate_size": ffn.intermediate_size}

Parameters:

block_config (BlockConfig)

Return type:

Dict[str, Any]

classmethod build_sequential_pipeline_module_fqns(*, num_stages, num_layers, layer_fqn_template, first_stage_fqns=(), last_stage_fqns=(), all_stage_fqns=())

Build an even sequential layer split for descriptor-defined PP layouts.

Parameters:

• num_stages (int)

• num_layers (int)

• layer_fqn_template (str)

• first_stage_fqns (Iterable[str])

• last_stage_fqns (Iterable[str])

• all_stage_fqns (Iterable[str])

Return type:

list[list[str]]

classmethod checkpoint_equivalence_tolerances()

Numerical output-space gates for function-preserving checkpoint transforms.

Most BF16 checkpoints use the strict defaults. Descriptors for storage formats whose kernels amplify otherwise equivalent permutations (for example MXFP4) may relax these output-space gates without weakening structural validation.

Return type:

dict[str, float]

classmethod checkpoint_key_candidates_for_model_key(model_key, *, model, config)

Return checkpoint keys that may satisfy a runtime model state_dict key.

Parameters:

• model_key (str)

• model (Module)

• config (Any)

Return type:

tuple[str, …]

classmethod create_dummy_block(original_layer, block_index)

Create a dummy block to replace a layer for sharded model initialization.

Parameters:

• original_layer (Module)

• block_index (int)

Return type:

Module

classmethod create_runtime_benchmark_model(runtime_config, block_configs)

Build a small model for vLLM latency benchmarking.

Implement this on descriptors that support runtime stats. Keeping model construction on the descriptor prevents the central benchmarking loop from hardcoding architecture-specific attention or MLP classes.

Parameters:

• runtime_config (Any)

• block_configs (list[BlockConfig])

Return type:

Module

abstract static decoder_layer_cls()

Decoder layer class types to patch for heterogeneous config support.

In most cases this class will hold as attributes both FFN & attention layers.

Returns:

nn.Module class type or a list if several class types should be patched.

Return type:

Type[Module] | List[Type[Module]]

classmethod embedding_pruning_spec(config, *, widths, alignment)

Return the complete residual-width contract, or reject unsupported families.

Parameters:

• widths (Iterable[int])

• alignment (int)

abstract static final_norm_name()

Return the name of the final normalization layer.

classmethod generic_decoder_contract(config)

Return this family’s composable decoder contract when one is declared.

Existing descriptors remain valid without adopting the generic surface. New cross-family descriptors override this hook so scoring, sorting, materialization, and preflight all consume the same structural declaration.

static get_language_model_config(config)

Get the language model config from a PretrainedConfig.

For regular LM models, returns the config itself. For VL/multimodal models with nested configs, override to return the language model portion (e.g., config.text_config for Qwen-VL).

classmethod get_passthrough_weight_groups(layer_names)

Group passthrough weights using passthrough_weight_name_predicates.

Parameters:

layer_names (Iterable[str])

Return type:

Dict[str, List[str]]

classmethod get_weight_groups(layer_names, num_hidden_layers)

Group model weights to support the puzzle subblock checkpointing format.

This method uses the abstract method layer_name_predicates by default.

Parameters:

• layer_names (Iterable[str]) – state_dict layer names of the model.

• num_hidden_layers (int) – number of decoder layers in the model.

Returns:

>>> {
...     "embedding": ["model.embed_tokens.weight"],
...     "lm_head": ["lm_head.weight", "model.norm.weight"],
...     "block_0_ffn": ["model.layers.0.mlp.down_proj", ...],
...     "block_0_attention": ["model.layers.0.self_attn.q_proj", ...],
... }

Return type:

Dictionary of group names to list of layer names per group, e.g.

abstract static init_rotary_embedding(model, runtime)

Re-initiate the rotary embeddings based on an existing model.

In puzzletron we initiate a sharded model by first creating a meta model then replacing to the actual device by loading the state_dict with the real weights.

Rotary embeddings frequencies are tensor buffers that are created dynamically during init and are not part of the model state_dict, so cannot be restored after a meta device initialization.

abstract static input_embedding_name()

Return the name of the input embedding layer.

classmethod is_passthrough_weight_name(name)

Return whether name belongs to a passthrough weight group.

Parameters:

name (str)

Return type:

bool

abstract static layer_block_name(index)

Return the name of the decoder layer at the given index.

Parameters:

index (int)

abstract static layer_name_predicates(num_layers)

Return predicates for grouping model weights to support subblock checkpointing.

For every group name return a regex predicate whether a layer name is part of the group.

Returns:

Dictionary of group name to regex pattern predicate.

Parameters:

num_layers (int)

Return type:

Dict[str, Pattern]

classmethod local_kd_subblock_module_paths(block_config, *, layer_idx)

Map semantic subblock identities to decoder-layer-relative module paths.

Parameters:

• block_config (BlockConfig)

• layer_idx (int)

Return type:

dict[tuple[str, str], str]

static mlp_no_op_post_init(decoder_layer)

Post-init callback to alter a decoder layer so that FFN/mlp subblock performs as no-op.

It is recommended to use the utils modules from no_op.py to replace layers to dummy counterparts.

Example for replacing a layernorm layer with identity:

>>> decoder_layer.post_attention_layernorm = Same()

Example for replacing an MLP layer with zeroes (zeroes since hidden_states are added to the residuals hidden_states so a no-op implementation will leave residual the same):

>>> decoder_layer.mlp = MatchingZeros()

In case the MLP layer to replace returns multiple outputs i.e hidden_states, _ = self.mlp(), use the util method return_tuple_of_size to return trailing None values:

>>> decoder_layer.mlp = return_tuple_of_size(MatchingZeros, size=2)()

Parameters:

decoder_layer (Module)

classmethod mlp_no_op_supported()

Check whether mlp_no_op_post_init is overridden for mlp no-op support.

Return type:

bool

abstract static output_embedding_name()

Return the name of the output embedding layer.

static passthrough_weight_name_predicates()

Return optional non-model weight groups that should be preserved as-is.

These tensors are not loaded into the active HF model but should survive conversion and checkpoint realization, e.g. draft/MTP heads ignored by the main model class.

Return type:

Dict[str, Pattern]

classmethod patch_layer_config(layer_config, block_config, layer_idx)

Apply structural per-layer fields not expressible as scalar overrides.

Most families need no work here. Hybrid families can update list-valued fields such as layer_types on the already-copied layer config.

Parameters:

• layer_config (Any)

• block_config (BlockConfig)

• layer_idx (int)

Return type:

None

classmethod patch_pipeline_model_part(model_part)

Patch a local AutoModel pipeline chunk after NeMo splits it.

This is a descriptor escape hatch for PP forward-name compatibility only. It should install transient aliases or lightweight attributes on the already-split stage object, not mutate canonical checkpoints or global library code. Return True when any patch was applied so callers can log what happened.

Parameters:

model_part (Module)

Return type:

bool

classmethod pipeline_module_fqns_per_model_part(config, *, pp_size, pipeline_config=None)

Return descriptor-owned pipeline stage module FQNs for NeMo AutoModel.

NeMo’s generic HF splitter assumes common names such as model.embed_tokens and model.norm. Remote-code families may use different module names while still being valid HF models; those names belong in the model descriptor so future families can customize PP splitting without patching NeMo or central Puzzletron code.

Return None to let NeMo AutoModel use its built-in split logic.

Parameters:

• config (Any)

• pp_size (int)

• pipeline_config (dict[str, Any] | None)

Return type:

list[list[str]] | None

classmethod ple_pruning_spec(config)

Return a global per-layer-embedding pruning contract when supported.

static position_id_axes(config)

Return the leading coordinate-axis count for canonical position IDs.

Standard decoder models use [batch, sequence] position IDs and therefore return one. Models whose distributed kernels require pre-expanded multi-axis positions (for example mRoPE) override this descriptor contract. Expansion happens before CP/PP sharding so those transforms preserve the coordinate axes.

Return type:

int

classmethod postprocess_runtime_benchmark_checkpoint(output_dir)

Descriptor hook for temporary vLLM benchmark checkpoint fixes.

Parameters:

output_dir (Any)

Return type:

None

static pruning_mixins()

Return available pruning mixins for bypass distillation.

Override in subclasses to provide model-specific pruning mixins, e.g. {"kv_heads": KVHeadsPruningMixIn(...), "experts_removal": ExpertRemovalPruningMixIn(...)}.

Returns an empty dict by default so that descriptors that do not need model-specific weight-slicing (e.g. Llama with standard FFN truncation) can rely on the generic create_child_state_dict fallback path.

Return type:

Dict[str, Any]

static requires_trust_remote_code()

Whether this model descriptor requires trust_remote_code=True for loading.

Models that use custom code (e.g., via auto_map in config) should override this to return True.

Returns:

True if trust_remote_code=True is required, False otherwise.

Return type:

bool

classmethod runtime_benchmark_base_block_config(runtime_config)

Return the standard block used as benchmark scaffolding.

Runtime stats measure a candidate subblock by repeating it after one standard block, then subtracting a matching baseline. Descriptors may override this for hybrid families whose default attention/MLP classes need extra config.

Parameters:

runtime_config (Any)

Return type:

BlockConfig

classmethod runtime_benchmark_config_fields(lm_config)

Return model-family fields required to synthesize latency benchmark configs.

Parameters:

lm_config (Any)

Return type:

dict[str, Any]

classmethod runtime_benchmark_export_descriptor()

Return the descriptor that matches the temporary benchmark checkpoint layout.

Return type:

type[ModelDescriptor]

classmethod runtime_benchmark_scaffold_policy(block_config)

Return the structural scaffold required for this runtime candidate.

Parameters:

block_config (BlockConfig)

Return type:

str

classmethod runtime_benchmark_sublayers_are_exclusive()

Whether each runtime layer executes exactly one active sublayer kind.

Return type:

bool

classmethod runtime_benchmark_supported()

Whether this descriptor implements synthetic vLLM benchmark models.

Return type:

bool

classmethod runtime_vllm_benchmark_args(config)

Return extra vllm bench latency args for this descriptor.

Parameters:

config (dict[str, Any])

Return type:

list[str]

classmethod set_block_configs(model_config, block_configs)

Attach block configs and update the language model layer count.

Multimodal configs often store the decoder layer count on a nested text config. This helper keeps callers from writing to config.num_hidden_layers directly when the real language model config lives elsewhere.

Parameters:

• model_config (Any)

• block_configs (list[BlockConfig | dict])

Return type:

None

classmethod split_passthrough_state_dict(state_dict)

Split a state dict into model weights and passthrough-only weights.

Parameters:

state_dict (dict[str, Any])

Return type:

tuple[dict[str, Any], dict[str, Any]]

classmethod stage_execution_policy()

Descriptor-owned execution exceptions for model stages.

Activation diagnosis repeatedly loads and evaluates physically different tensor shapes in one process. Globally compiled model-local kernels can retain shape-specific distributed graphs across those transitions, so diagnosis defaults to the mathematically identical eager path. A descriptor may extend or override this policy. Keeping it descriptor- owned avoids model-name checks in launchers and lets future families reuse the same generic mechanism.

Return type:

dict[str, tuple[str, …]]

static truncate_pattern_for_subblock(lm_config, parent_layer_index=None)

Adjust per-layer config fields so a single-layer model represents the correct layer type.

The default implementation handles hybrid_override_pattern for hybrid architectures. It is a no-op when the field is absent. Override if a model uses a different pattern alphabet.

Parameters:

• lm_config (Any)

• parent_layer_index (int | None)

Return type:

None

classmethod update_runtime_benchmark_config(config_data)

Adjust the temporary benchmark config before vLLM loads it.

Parameters:

config_data (dict[str, Any])

Return type:

None

static uses_autocast()

Whether this model supports torch.autocast.

Some models (e.g., Qwen3-VL MoE) have dtype bugs under autocast. Override and return False for models that do not support autocast.

Return type:

bool

classmethod vision_module_names()

Canonical vision-tower module names for multimodal observability.

Return type:

tuple[str, …]

classmethod width_slice_equivalence_operations(config, sorted_checkpoint_dir, *, alignment=1, sampled_layers=None)

Build declarative cases for every capability-backed width operation.

Parameters:

• alignment (int)

• sampled_layers (Iterable[int] | None)

classmethod width_slice_equivalence_tolerances()

Numerical gates for physical-versus-runtime width slices.

Return type:

dict[str, float]