Nested Local Distillation for Puzzletron

Feature Request: Nested Local Distillation for Puzzletron

Description

The current evaluation and training setup requires running either N (number of FFN options) + M (number of attention options) separate bypass runs, or in worst-case scenarios, a full combinatorial grid of N * M runs. As target architectures grow more intricate, this configuration space scales exponentially and quickly becomes completely intractable.

For example, when optimizing a highly complex model like nemotron3-super, the FFN option space (N) alone is composed of a massive combination of parameters:

Options for N = num_experts_options (~12) * expert_channels_options (~8) * latent_dim_options (~4) * top_k_options (~4) * shared_expert_dim_options (~8) ~= 12,288

Evaluating this combinatorial explosion via independent bypass runs creates an enormous computational bottleneck. To solve this, we propose adding a Nested Local Distillation capability as an optional feature within Puzzletron.

Implementation Concept

Instead of launching distinct distillation runs for every single configuration variant, the proposed feature integrates a dynamic, single-run optimization strategy:

• Stochastic Configuration Sampling: At the start of each minibatch during the distillation phase, the pipeline dynamically samples a valid model configuration from the architectural search space.
• Matryoshka-Style Weight Selection: The system utilizes nested, Matryoshka-style weight selection to slice out sub-networks corresponding to the sampled configuration.
• Unified Weight Updates: Sub-network parameters are optimized and updated concurrently within the single shared run, ensuring that smaller nested weights are systematically optimized alongside the larger, overarching representations.

Expected Impact

Introducing Nested Local Distillation will collapse the exponential N * M run matrix into a single, highly efficient profiling session. This dramatically reduces the compute budget and wall-clock time required to compress next-generation Mixture of Experts (MoE) and hybrid architectures, unlocking scalability for massive foundational checkpoints.

Furthermore, this unified, pre-sorted weight structure serves as an excellent initialization point for Flextron, streamlining downstream continuous-elasticity training workflows without requiring heavy reconfiguration or cold-start warming.