automodel#

Native NeMo AutoModel bridge for Puzzletron AnyModel checkpoints.

Classes

AutoModelDescriptor

Descriptor for NeMo AutoModel native blocks.

AutoModelDescriptorFactory

ContractAutoModelDescriptor

Native bridge that delegates shape semantics to an AnyModel contract.

Functions

automodel_patcher

Temporarily patch native AutoModel decoder block construction.
class AutoModelDescriptor#

Bases: object

Descriptor for NeMo AutoModel native blocks.

AutoModel bypass works by making teacher and student native AutoModel builds enter the same heterogeneity context. The patcher temporarily wraps each decoder block __init__ so layer i receives overrides derived from block_configs[i] before parameters are allocated; after construction the descriptor can replace no-op submodules and patch state-dict adapters. Local bypass then captures descriptor-declared block/subblock I/O on the native model parts and trains only the candidate submodule, while global KD uses the same descriptor to keep teacher/student sharding and per-layer shapes aligned.

static attn_no_op_post_init(layer)#
Parameters:

layer (Any)

Return type:

None

static block_config_to_config_overrides(block_config)#
Parameters:

block_config (BlockConfig)

Return type:

dict[str, Any]

static decoder_layer_cls()#
Return type:

type | tuple[type, …]

classmethod make_patched_init(orig_init, block_configs)#
Parameters:
Return type:

Callable

static mlp_no_op_post_init(layer)#
Parameters:

layer (Any)

Return type:

None

classmethod native_state_dict_adapter_context(block_configs)#

Temporarily adapt a native backend’s checkpoint shape bridge.

classmethod patch_constructor_arguments(arguments, block_config, layer_idx)#

Patch non-config constructor inputs for a native backend.

Most native blocks derive all geometry from config. Backends that pass a separate structural object (for example a sharded MoE config) override this hook while retaining the common signature-aware patcher.

Parameters:

  • arguments (dict[str, Any])

  • block_config (BlockConfig)

  • layer_idx (int)

Return type:

None

static patch_hf_model_checkpoint_mapping(model)#
Parameters:

model (Any)

Return type:

bool

classmethod patch_layer_config(config, block_config, layer_idx)#

Apply structural per-layer fields shared by native model families.

Scalar projection sizes are handled by block_config_to_config_overrides. Windowed attention also changes a layer’s attention kind, so it belongs here and is applied to the private config copy passed to that layer.

Parameters:

  • config (Any)

  • block_config (BlockConfig)

  • layer_idx (int)

Return type:

None

static patch_state_dict_adapter(model)#
Parameters:

model (Any)

Return type:

bool

class AutoModelDescriptorFactory#

Bases: object

classmethod get(name)#
Parameters:

name (str)

Return type:

type[AutoModelDescriptor] | None

classmethod register(**entries)#
Parameters:

entries (type[AutoModelDescriptor])

Return type:

None

classmethod register_decorator(*names)#
Parameters:

names (str)

Return type:

Callable

classmethod registered_names()#
Return type:

tuple[str, …]

class ContractAutoModelDescriptor#

Bases: AutoModelDescriptor

Native bridge that delegates shape semantics to an AnyModel contract.

A family-specific subclass only declares the corresponding structural descriptor and native block class. Scoring, HF construction, native construction, materialization, and vLLM then share one field mapping.

STRUCTURAL_DESCRIPTOR: Any = None#
classmethod block_config_to_config_overrides(block_config)#
Parameters:

block_config (BlockConfig)

Return type:

dict[str, Any]

classmethod patch_layer_config(config, block_config, layer_idx)#
Parameters:

  • config (Any)

  • block_config (BlockConfig)

  • layer_idx (int)

Return type:

None

automodel_patcher(descriptor, block_configs)#

Temporarily patch native AutoModel decoder block construction.

Parameters: