Add stochastic mixer for supernet training

[tscholak]

Summary

Implements a stochastic mixer layer for supernet training, enabling random sampling from multiple mixer options (e.g., attention vs. Mamba) during training. Includes checkpoint conversion support and a hierarchical beam search tool for finding optimal mixer placement post-training.

Implementation Details

Stochastic Mixer (fast_llm/layers/decoder/stochastic_mixer.py)

• Training mode: Randomly samples from configured mixers using distributed RNG (ensures consistency across TP/PP ranks)
• Eval/inference mode: Uses main_mixer_index for deterministic behavior
• Sampling strategies: Uniform or weighted sampling
• Preprocessing: Runs preprocessing for all mixers since we don't know which will be selected

Configuration (fast_llm/layers/decoder/config.py)

• StochasticMixerConfig: List-based mixer configuration with sampling strategy
• main_mixer_index: Specifies which mixer to use during inference and which receives pretrained weights during checkpoint conversion
• Validation ensures sampling weights sum to 1.0 and all indices are valid

Checkpoint Conversion (fast_llm/models/gpt/conversion/apriel.py)

• AprielStochasticMixerConverter: Handles conversion between Fast-LLM and Apriel formats
• Only main_mixer_index weights are exported/imported (other mixers randomly initialized during supernet training)
• Follows existing converter patterns (minimal, no verbose comments)

Beam Search Tool (tools/supernet_beam_search.py)

• Hierarchical algorithm: Finds optimal placement for mixers at each quality/cost level
  • Phase 1: Find best N layers for primary mixer (e.g., full attention)
  • Phase 2: Find best M layers for secondary mixer (e.g., sliding window attention)
  • Remaining layers use tertiary mixer (e.g., linear attention)
• Efficient evaluation: Loads checkpoint once, modifies main_mixer_index in-place for each candidate
• Fast-LLM integration: Uses Fast-LLM's evaluation system directly (no subprocess or checkpoint reconversion)
• Features: Pre-scoring, beam growth, early stopping, configurable score direction

Tests (tests/utils/model_configs.py)

• Added stochastic_mixer test configuration with FA/Mamba mixers
• Enabled checkpoint conversion testing via AprielHybridSSMCheckpointFormat

Use Case

Supernet Training: Train a model where each layer can be either full attention or Mamba, with random sampling at each step. After training, use beam search to find which specific layers benefit most from full attention vs. Mamba, given a budget constraint (e.g., "I can afford 4 FA layers").

Testing

Run the stochastic mixer tests:

pytest tests/models/test_checkpoint.py::test_checkpoint_and_eval tests/models/test_checkpoint.py::test_conversion -k "stochastic_mixer" -v

Example beam search usage:

fast-llm tools/supernet_beam_search.py \
  training_config=path/to/supernet_config.yaml \
  budgets=[4,8] \
  beam_width=12 \
  score_metric="lm_eval/accuracy" \
  output_path=results.json

🤖 Generated with Claude Code

Co-Authored-By: Claude [email protected]

[jlamypoirier]

Looks good, some minor comments

[jlamypoirier]

This requires a costly cuda sync. How about we sample for all layers at once during preprocessing?

[tscholak]

now done during preprocessing

[jlamypoirier]

Ambiguous if multiple mixers share the same type. Use named mixers instead?

[tscholak]

now using named mixers. we retrieve mixer_name from kwargs (line 151) and use it for logging (line 160) and accessing the correct mixer (line 163).

[jlamypoirier]

There could be name conflicts. Consider namespace?

[tscholak]

now namespaced. see lines 214-216 where we prefix with f"{mixer_name}/{loss_def.name}".

[jlamypoirier]

This is a bit dangerous, there could be name conflicts and counts will be wrong for averaging. Not sure how to fix though.

[tscholak]

Acknowledged. The current approach ensures we allocate space for all possible losses, but you're right that counts won't match actual usage since only one mixer runs per forward pass. We could track which mixer was use and only record its losses, but that adds complexity. I think what we have is good enough for now.

[jlamypoirier]

hf_prefix. drop_on_export handles the rest.

[tscholak]

now uses just hf_prefix without the mixer name prefix.

[jlamypoirier]

I don't think match is warranted here, since it involves a (slow) initialization of configs.

[tscholak]

uses if-else with instance type check now.

[jlamypoirier]

We could just enforce this and make main_mixer_name a cached property instead?

[tscholak]

I like it the way it is

[jlamypoirier]

Could we just use a cpu generator instead? Or at the very least move to device in setup

[jlamypoirier]

I don't think that's right, tp_generator will result in different tensor ranks selecting different mixers.

[jlamypoirier]

Also do we actually want different DP ranks / micro-batches to select different sets? I guess this increases randomness but it will affect reproducibility and prevent distributed tests from working.

[tscholak]

thanks for checking this. what we want/need is that all ranks sample the same mixer for each batch. How can that be done? I thought that's what the tp generator does. maybe it does the exact opposite, and all tp ranks do it differently?

[jlamypoirier]

The TP generator is meant for TP tensors, which need different random numbers for each TP ranks. (Ex. for we want different dropouts for different slices of a tensor).

I had a second look, and I don't think any existing generator can provide consistent mixers for a given batch. pp_generator gives inconsistent results between DP ranks and gradient accumulation steps, but is probably still the best option. The CPU generator on the other hand is consistent between DP ranks (not grad accumulation), but is not reproducible (see Distributed.set_step). Getting consistency for a whole batch would require a custom seed/generator, and access to the current training step in preprocess. (I guess preprocess_batch could add it to the kwargs).

Another issue I'm seeing, preprocess is called only once for all layers so I think the current approach will result in all layer choosing the same mixer. And fixing is non-trivial since the preprocessor doesn't know the number of layers and the mixer doesn't know about its layer index. So my suggestion would be to go with a custom generator, seed it in preprocess using the step index, then generate on the fly in forward

[tscholak]

Hi @jlamypoirier! Thanks for the thorough review! I've addressed all your feedback and cleaned up the
PR significantly. Here's what changed since your last review:

Stochastic mixer sampling fix (your main concern):

• Added iteration to BlockKwargs and propagated it through preprocess_batch → kwargs
• Create a torch CPU generator in preprocess(), seeded with iteration
• Sample mixers in forward() using torch.multinomial with the CPU generator
• This ensures all DP/TP/PP ranks sample the same mixer sequence per batch, while different layers
  sample different mixers deterministically
• Removed the TP/PP generator usage that was causing rank inconsistencies
• Store sampling probabilities on CPU to avoid device transfers

Code cleanup:

• Removed all debug print statements and irrelevant changes
• Removed stochastic mixer support from old apriel converter (only supported in apriel2 now)
• Removed tools/supernet_beam_search.py (will implement differently later)
• Renamed SamplingStrategy → StochasticMixerSamplingStrategy
• Use normalize_probabilities for sampling weight validation (consistent with dataset blending)

Note: I added a new apriel2 conversion module for the newer model format. This doesn't need review as it's separate from the stochastic mixer changes. Apriel hybrid SSM export is not supported. Stochastic layers are only in apriel2.

Ready for another look!