base

Classes

Converter

Base class for adding AnyModel metadata to HuggingFace checkpoints.
class Converter

Bases: ABC

Base class for adding AnyModel metadata to HuggingFace checkpoints.

classmethod convert(descriptor, input_dir, output_dir)

Attach AnyModel block configs to a standard HuggingFace checkpoint.

The normal conversion path copies the original HF checkpoint files, including safetensors and index files, then writes config artifacts with typed block_configs. It does not reorganize model weights.

Parameters:

  • descriptor (ModelDescriptor) – Model descriptor for the model type.

  • input_dir (Path) – Path to the input HuggingFace checkpoint.

  • output_dir (Path) – Path to the output AnyModel checkpoint.

classmethod convert_configs_in_dirs(input_dir, output_dir, trust_remote_code=False, descriptor=None)

Attach typed block_configs and save standard HuggingFace config artifacts.

Parameters:

  • input_dir (Path)

  • output_dir (Path)

  • trust_remote_code (bool)

  • descriptor (ModelDescriptor | None)

static copy_checkpoint_files(input_dir, output_dir)

Materialize a checkpoint while preserving standard HuggingFace weight files.

Metadata/config files are copied so Puzzletron can rewrite config.json without mutating the source checkpoint. Large weight shards are hardlinked when possible and copied otherwise. This keeps 100B+ smoke conversions fast when the HF cache and run directory share a filesystem, while still producing self-contained checkpoints when they do not.

Parameters:

  • input_dir (Path)

  • output_dir (Path)

abstract static create_block_configs_from_main_config(config)

Create per-layer BlockConfig list from a HuggingFace model config.

This method extracts layer-specific parameters (e.g., intermediate_size, num_kv_heads) from the main model config and creates a BlockConfig for each layer. These BlockConfigs enable layer-specific pruning and modifications during the compression pipeline.

Parameters:

config (PreTrainedConfig) – HuggingFace PretrainedConfig (e.g., LlamaConfig, Qwen2Config)

Returns:

List of BlockConfig, one per hidden layer. Each BlockConfig contains typed subblock_configs, for example AttentionConfig, FFNConfig, MoEConfig, or MambaConfig entries.

Return type:

List[BlockConfig]

Example

For a model with uniform layers (e.g., Llama):

return [BlockConfig(…)] * config.num_hidden_layers

For a model with heterogeneous layers (e.g., NemotronH with Mamba/Attention):

return [BlockConfig(…) for layer_idx in range(num_layers)]