huggingface / huggingface/candle
qwen3_5: the sparse-MoE layers and partial rotary the released checkpoints need
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Problem
Both Qwen 3.5 implementations in this crate — `qwen3_5.rs` and
`quantized_qwen3_5.rs` — are the **dense** variant. Neither mentions experts:
```
$ grep -ic expert candle-transformers/src/models/qwen3_5.rs
0
$ grep -ic expert candle-transformers/src/models/quantized_qwen3_5.rs
0
```
`Qwen3_5MLP` is a plain SwiGLU over `intermediate_size`, and `TextConfig` has
no `num_experts`, `num_experts_per_tok`, `moe_intermediate_size`,
`shared_expert_intermediate_size`, or `partial_rotary_factor`.
The published Qwen 3.5 checkpoints are mixture-of-experts with a shared expert.
So the model in this file is not the model the weights describe, and the two
config gaps fail in different ways:
- **The MoE gap fails loudly.** A sparse layer stores
`mlp.experts.{i}.{gate,up,down}_proj`, `mlp.gate` and `mlp.shared_expert.*`;
`Qwen3_5MLP::new` asks for `mlp.gate_proj.weight` and gets a missing-tensor
error. Unpleasant, but honest.
- **`partial_rotary_factor` fails silently.** `TextConfig` derives plain
`Deserialize` with no `deny_unknown_fields`, so the field is dropped on the
floor, and `Qwen3_5TextRotaryEmbedding::new` builds its table over
`dim = cfg.head_dim()` — the whole head. A checkpoint that rotates only the
first `head_dim * partial_rotary_factor` dimensions loads without a word and
produces wrong numbers. This is the same shape of problem as the
`attention_bias` hole fixed in #3832: the config says something the model
never reads.
## Why this is worth closing now
#3832 gave the layers a projection seam and #3831 gave `candle-nvfp4-kernels` a
device-resident `Module`. Between them, a caller with a ModelOpt NVFP4
checkpoint can supply its own projections and run everything else — attention,
linear attention, `gated_delta_rule_chunked`, norms, cache, residuals — out of
this crate.
Everything except the layer the checkpoint actually has. So the seam currently
opens onto a model those checkpoints cannot be loaded into, and a caller is
back to reimplementing the file it was meant to reuse. Closing this gap is what
turns those two changes from a good API into a working path.
## What the sparse layer needs
The semantics below are the ones we run in production against real Qwen 3.5
NVFP4 checkpoints, with a scalar reference implementation and its tests.
1. **Router top-k over renormalised probabilities.** Softmax the router logits,
take the top `num_experts_per_tok`, then renormalise the selected weights so
they sum to one. Ties broken by ascending expert index, so a run is
reproducible.
2. **A shared expert on every token**, over `shared_expert_intermediate_size`,
scaled by `sigmoid` of the `mlp.shared_expert_gate` logit and added to the
routed sum. It is not one of the `num_experts`.
3. **Routed experts over `moe_intermediate_size`**, which differs from the
dense `intermediate_size` and is what sizes the expert weights.
4. **Grouping tokens by expert, not looping tokens.** Each expert's weights are
read once for all the tokens routed to it in a chunk. Done per token, an
expert's operator is re-read for every activation, and on a quantized
checkpoint that means re-decoding its weights — which for us was the single
largest cost in a prefill.
And for the rotary: `rotary_dim = (head_dim as f64 * partial_rotary_factor) as
usize`, rotating the leading `rotary_dim` values of each head and passing the
rest through. A missing field should default to `1.0` so dense checkpoints are
unaffected.
## Shape
Following `layer_types`, which already selects the token mixer per layer, the
natural form is a second arm on the same kind of switch:
```rust
enum FeedForward {
Dense(Qwen3_5MLP),
Sparse(Qwen3_5SparseMoeBlock),
}
```
with the projection seam from #3832 reaching the experts by name
(`experts.{i}.gate_proj`, `shared_expert.up_proj`, `gate`,
`shared_expert_gate`) exactly as it reaches the dense ones — so a quantized
checkpoint gets its experts through the same factory, with no new API.
Alternatively, a `BlockMlp`-style trait as `llama.rs` has (#3809) would let a
caller supply the whole mixture. I would rather have the mixture in the crate:
the routing rule, the renormalisation and the shared-expert gate are model
semantics, and every caller reimplementing them is how they drift.
## Note
Happy to submit the PR. The tests I would bring are the two that caught real
bugs for us: a batched mixture must equal running the same tokens one at a
time, bit for bit, so chunk size stays a performance knob; and each expert must
be invoked once per chunk rather than once per token, which is the property the
grouping exists for and the one a refactor silently loses.
Related: #3809, #3831, #3832.
Contributor guide
No contributing guide indexed for this repository
Research direction
Read candle-transformers/src/models/qwen3_5.rs and quantized_qwen3_5.rs, then compare the layer_types switch and projection seam from #3832. Verify the sparse implementation against the stated router and shared-expert semantics, including partial rotary handling; done means released checkpoints load and batched results match token-at-a-time execution while each expert is invoked once per chunk.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 45/100