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)

Open
#71 1 comment 0 reactions 0 assignees View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.