Lazy Checkpoint Creation for Puzzletron

Feature Request: Lazy Checkpoint Creation for Puzzletron

Description

The current workflow for including pruning configurations into the Puzzletron search space requires creating a dedicated, homogeneously pruned checkpoint for every single target ratio. While this is manageable for small models or limited combinatorial search spaces, it becomes an unsustainable bottleneck for storage and I/O as model sizes and architectural complexities scale.

To resolve this, we should adopt a lazy checkpoint creation strategy. The core proposal is to sort the checkpoint weights exactly once (similar to the methodology used in Flextron). For all subsequent stages of the pipeline, instead of reading from separate pre-pruned files on disk, we will dynamically slice the channels in-memory while loading the model. This single architectural shift will drastically reduce storage footprints by a significant factor.

Pipeline Verification Requirements

Implementing this feature requires verifying that all downstream stages of the Puzzletron pipeline can seamlessly ingest and operate on these dynamically sliced, sorted checkpoints. Specifically, we must audit and adapt the following four core stages:

1. Bypass: Ensure that the bypass mechanisms correctly map the sliced dimensions without expecting explicit on-disk structure changes.
2. Stats / vLLM: Verify that weight slicing integrates with statistics collection and remains compatible with vLLM's block discovery.
3. Replace-1-Block-Scoring: Validate that block-by-block replacement and scoring correctly compute importance metrics using the dynamically sliced tensors.
4. Model Realization after MIP: Confirm that the final model realization step can successfully materialize the optimal architecture found by the Mixed-Integer Programming (MIP) solver, writing out the final, fully pruned checkpoint to disk.

Expected Impact

Transitioning to lazy checkpoint creation will eliminate the exponential storage overhead associated with dense combinatorial search spaces. This allows the team to evaluate significantly larger models and wider search granularities without running into infrastructure or disk space constraints.