elastic_sampling

Prefix-size sampling for nested (Matryoshka) elastic bypass distillation.

Each prunable subblock samples one of its allowed sizes per minibatch with probability p ~ 1 / num_params(size) (normalized). Larger variants (more params, closer to the teacher) are sampled less, so they get fewer recovery iterations; small variants get more. The full (teacher) size is included.

The parameter counts MUST come from the canonical subblock_stats.calc_subblock_params_and_memory.calculate_subblock_params (which builds the meta decoder layer and counts sum(p.numel()), so GQA, gated-vs-not FFN, biases, and norms are all exact) — not a hand-written formula. This module only consumes the counts (so it stays pure / unit-testable); elastic_supernet.build_subblock_elastics supplies them via that function.

Classes

ElasticSizeSampler

Samples a size for one subblock with p ~ 1/num_params (full/teacher size included).

Functions

inverse_param_probs	Normalized probabilities proportional to 1 / param_count (float64).
inverse_width_probs	Normalized hidden-width probabilities proportional to 1 / width.

class ElasticSizeSampler

Bases: object

Samples a size for one subblock with p ~ 1/num_params (full/teacher size included).

sizes is the allowed list (FFN intermediate ints, or (q, kv) tuples for attention); param_counts are the corresponding parameter counts. Use one sampler per prunable subblock. A shared torch.Generator makes a run reproducible.

__init__(sizes, param_counts)

sample(generator=None)

Parameters:

generator (Generator | None)

inverse_param_probs(param_counts)

Normalized probabilities proportional to 1 / param_count (float64).

No-op elastic endpoints have zero parameters. Treat them as having the same effective size as the smallest non-zero candidate so they are sampled as a normal endpoint instead of dominating the distribution or producing inf.

Return type:

Tensor

inverse_width_probs(widths)

Normalized hidden-width probabilities proportional to 1 / width.

Return type:

Tensor