Add stochastic mixer for supernet training#373

Open

tscholak wants to merge 28 commits intomainfrom

stochastic-mixer

Commits on Oct 12, 2025

Add stochastic mixer for supernet training

tscholak
and
claude
committed

Commits on Oct 14, 2025

Fix stochastic mixer test failures

tscholak
and
claude
committed
Fix stochastic mixer checkpoint conversion

tscholak
and
claude
committed

Commits on Nov 15, 2025

wip
tscholak
committed

Commits on Nov 20, 2025

Commits on Nov 21, 2025

Add Apriel2 checkpoint format and fix weight tying

tscholak
and
claude
committed
Optimize Apriel2: compute position embeddings and masks per unique block

tscholak
and
claude
committed
Add HuggingFace generation and caching improvements to Apriel2

tscholak
and
claude
committed
Add Apriel2DynamicCache for hybrid attention/SSM layer support

tscholak
and
claude
committed
Add Mamba incremental generation support to Apriel2

tscholak
and
claude
committed
Add GatedDeltaNet support via Qwen3NextGatedDeltaNet wrapper

tscholak
and
claude
committed
Standardize naming: recurrent_states and Apriel2 prefixes

tscholak
and
claude
committed
Remove debug print statements and irrelevant changes

tscholak
and
claude
committed
Remove stochastic mixer support from apriel conversion

tscholak
and
claude
committed
Remove trivial formatting change from apriel_hybrid_ssm config

tscholak
and
claude
committed
Remove test changes for lossy HF conversion

tscholak
and
claude
committed
Revert trivial setup.py formatting and restore .eval() calls in tests

tscholak
and
claude
committed
Rename SamplingStrategy to StochasticMixerSamplingStrategy

tscholak
and
claude
committed
Use normalize_probabilities for sampling weights validation

tscholak
and
claude
committed
Remove tools/supernet_beam_search.py

tscholak
and
claude
committed

Commits on Nov 22, 2025