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: