deepseek-ai / deepseek-ai/DeepSpec
DSpark Qwen3 draft modeling silently ignores `partial_rotary_factor`, producing checkpoints whose config mismatches the trained weights (breaks vLLM serving AL)
- Dominant language
- Python
- Stars
- 7.1k
- Forks
- 667
- PR merge metrics
- No merged PRs in 30d
Description
### Summary
The DSpark Qwen3 draft modeling silently ignores `partial_rotary_factor` declared in the draft config. Training therefore applies **full-head-dim RoPE**, but the field is saved back into the checkpoint's `config.json` verbatim. Serving engines that honor the field (e.g. vLLM's dflash/dspark path, merged in vllm-project/vllm#46995) then rotate only `head_dim * factor` dims — a train/serve attention-geometry mismatch that collapses acceptance length and is very hard to debug.
### Where
`deepspec/modeling/dspark/qwen3/modeling.py`:
- L34-L40 — the custom `apply_rotary_pos_emb` applies `cos`/`sin` across the **full** head dim, with no rotary/pass-through split:
```python
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
...
q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
k_embed = (k * cos) + (rotate_half(k) * sin)
```
- L239 — `self.rotary_emb = Qwen3RotaryEmbedding(config)`. With transformers 5.13.1 and a config declaring `partial_rotary_factor: 0.25`, `head_dim: 128`, this still yields **128-dim** cos/sin (`inv_freq` len 64), so the declared factor never takes effect at training time either way.
### How you hit this in practice
Draft configs for Qwen3.5-family targets naturally inherit `partial_rotary_factor: 0.25` from the target's attention geometry — e.g. z-lab's public DFlash draft heads ship with it, and they are the obvious warm-start initialization for DSpark training. After fine-tuning with DeepSpec, the weights are (re-)adapted to full-dim rotation, but `save_pretrained` writes the config back with `0.25`. The checkpoint is now self-inconsistent: weights say "full-dim RoPE", config says "rotate 32/128 dims".
vLLM's dspark/dflash serving path faithfully executes the config (`rotary_dim = 32`, affecting both the query-block forward and the per-layer context-KV precompute), so every draft backbone layer computes attention with different geometry than training.
### Observed impact (Qwen3.5-27B-class target, 5-layer DSpark draft, block 8, warm-started from z-lab DFlash)
- Offline DeepSpec eval: AL ≈ 2.53. vLLM serving with the same checkpoint: AL ≈ 1.4.
- Debugging is nasty because the target-side aux features match across engines to cos ≥ 0.9999 — only the draft backbone hidden states diverge.
- After setting `partial_rotary_factor: 1.0` in the checkpoint config (one-line fix, no engine changes): vLLM's per-row backbone hidden matches the DeepSpec replay at cos ≥ 0.99994, sampled draft chains are identical token-for-token, and serving AL recovers to 2.59 ≈ offline.
### Minimal check
```python
from transformers import AutoConfig
from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
import torch
c = AutoConfig.from_pretrained(CKPT) # declares partial_rotary_factor = 0.25
r = Qwen3RotaryEmbedding(c)
cos, _ = r(torch.zeros(1, 8, c.hidden_size), torch.arange(8)[None])
print(cos.shape) # [..., 128] -> training rotates all dims
```
vs. vLLM: `get_rope(128, max_position=..., rope_parameters=c.rope_parameters).rotary_dim` → `32`.
### Suggested fix
Any of these would prevent the silent mismatch:
1. Honor `partial_rotary_factor` in the DSpark/DFlash draft attention (split rotary/pass-through dims in `apply_rotary_pos_emb`), or
2. Normalize the field at draft-config build/save time (set to `1.0` or drop it) with a warning, or
3. At minimum, assert/warn at model init when the config declares a factor the modeling does not implement.
Happy to provide more details or a PR for option 2/3 if useful.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.