Skip to content

DSpark Qwen3 draft modeling silently ignores partial_rotary_factor, producing checkpoints whose config mismatches the trained weights (breaks vLLM serving AL) #71

Description

@arthurgao2003

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:
    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

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_dim32.

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions