build_replacement_library

This module constructs the replacement library JSON files from a puzzle directory containing multiple trained model checkpoints. It analyzes checkpoints to extract unique block and subblock configurations, builds a library of available replacements, and generates solutions for layer replacement in compressed models. The resulting replacement library can then be used by ReplacementLibrary to efficiently load models with mixed teacher/student layers.

Standard Puzzle Usage:

python -m modelopt.torch.puzzletron.replacement_library.build_replacement_library PUZZLE_DIR

Teacher checkpoint dir is assumed to be inside PUZZLE_DIR/ckpts/teacher (symlink is recommended) though you can supply an explicit –teacher_checkpoint_dir.

–add_ffn_no_ops and –add_attention_no_ops are optional (default True),

Functions

build_replacement_library	For normal puzzle runs, use default values.
launch_build_replacement_library	Launch the build replacement library function with Hydra configuration.
infer_teacher_dir

build_replacement_library(master_puzzle_dir, descriptor, teacher_checkpoint_dir=None, add_ffn_no_ops=True, add_attention_no_ops=True)

For normal puzzle runs, use default values. For advanced use cases, see the Usage section.

Parameters:

• master_puzzle_dir (Path | str)

• descriptor (ModelDescriptor)

• teacher_checkpoint_dir (Path | str | None)

• add_ffn_no_ops (bool)

• add_attention_no_ops (bool)

Return type:

None

infer_teacher_dir(master_puzzle_dir, teacher_checkpoint_dir=None)

Parameters:

• master_puzzle_dir (Path | str)

• teacher_checkpoint_dir (Path | str | None)

Return type:

Path

launch_build_replacement_library(cfg)

Launch the build replacement library function with Hydra configuration.

Parameters:

cfg (DictConfig)

Return type:

None