From b66c2f372ad8efa9ada77ed470e51027518bb76d Mon Sep 17 00:00:00 2001 From: arthurgao2003 Date: Wed, 22 Jul 2026 16:24:28 +0800 Subject: [PATCH] Warn on and normalize partial_rotary_factor declarations DSpark does not implement DSpark's Qwen3 draft modeling applies RoPE across the full head dim, but build_draft_config deepcopies the target config, so a target (or DFlash warm-start checkpoint) declaring partial_rotary_factor != 1.0 leaks the field into saved draft checkpoints verbatim. Serving engines that honor it (e.g. vLLM's dflash/dspark path) then rotate fewer dims than training did, silently collapsing acceptance length (observed 2.53 offline -> ~1.4 in vLLM; recovers to 2.59 once the config declares 1.0). - build_draft_config now normalizes the declaration to 1.0 with a warning - Qwen3DSparkModel init warns when loading a config that still declares partial RoPE (covers pre-existing checkpoints) Fixes #71 Co-Authored-By: Claude Fable 5 --- deepspec/modeling/dspark/common.py | 49 ++++++++++++++++++++++ deepspec/modeling/dspark/qwen3/config.py | 9 +++- deepspec/modeling/dspark/qwen3/modeling.py | 5 +++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/deepspec/modeling/dspark/common.py b/deepspec/modeling/dspark/common.py index 76d908a4..f833697e 100644 --- a/deepspec/modeling/dspark/common.py +++ b/deepspec/modeling/dspark/common.py @@ -1,3 +1,4 @@ +import warnings from dataclasses import dataclass from typing import Optional @@ -56,6 +57,54 @@ def extract_context_feature(hidden_states, layer_ids): ) +def declared_partial_rotary_factor(config) -> float: + for source in ( + getattr(config, "rope_parameters", None), + getattr(config, "rope_scaling", None), + ): + if isinstance(source, dict) and source.get("partial_rotary_factor") is not None: + return float(source["partial_rotary_factor"]) + value = getattr(config, "partial_rotary_factor", None) + return 1.0 if value is None else float(value) + + +def check_partial_rotary_factor(config, normalize: bool = False): + """DSpark draft attention applies RoPE across the full head dim, so a config + declaring partial_rotary_factor != 1.0 (e.g. inherited from a Qwen3.5-style + target via copy.deepcopy, or from a DFlash warm-start checkpoint) is not what + training actually does. Serving engines that honor the field (e.g. vLLM's + dflash/dspark path) would then rotate fewer dims than training did, silently + collapsing the acceptance length. Warn, and optionally normalize the + declaration to 1.0 so saved checkpoints describe the trained geometry.""" + declared = declared_partial_rotary_factor(config) + if declared == 1.0: + return config + if normalize: + for source in ( + getattr(config, "rope_parameters", None), + getattr(config, "rope_scaling", None), + ): + if isinstance(source, dict) and "partial_rotary_factor" in source: + source["partial_rotary_factor"] = 1.0 + if getattr(config, "partial_rotary_factor", None) is not None: + config.partial_rotary_factor = 1.0 + warnings.warn( + f"Draft config declared partial_rotary_factor={declared}, but DSpark " + "modeling applies RoPE across the full head dim; normalized the " + "declaration to 1.0 so the saved checkpoint matches the trained " + "geometry." + ) + else: + warnings.warn( + f"Config declares partial_rotary_factor={declared}, but DSpark " + "modeling applies RoPE across the full head dim and ignores it. " + "Serving engines that honor this field (e.g. vLLM) will apply " + "different attention geometry than training; set it to 1.0 in the " + "checkpoint config before deploying." + ) + return config + + def validate_target_layer_ids(layer_ids, num_target_layers: int): layer_ids = [int(layer_id) for layer_id in layer_ids] assert layer_ids, "target_layer_ids must not be empty." diff --git a/deepspec/modeling/dspark/qwen3/config.py b/deepspec/modeling/dspark/qwen3/config.py index ae71facc..ee09259b 100644 --- a/deepspec/modeling/dspark/qwen3/config.py +++ b/deepspec/modeling/dspark/qwen3/config.py @@ -1,6 +1,9 @@ import copy -from deepspec.modeling.dspark.common import validate_target_layer_ids +from deepspec.modeling.dspark.common import ( + check_partial_rotary_factor, + validate_target_layer_ids, +) TRAIN_ATTN_IMPLEMENTATION = "flex_attention" @@ -35,6 +38,10 @@ def build_draft_config( ) draft_config = copy.deepcopy(target_config) + # The target config may declare partial RoPE (e.g. Qwen3.5 targets), which + # the DSpark draft modeling does not implement; normalize the declaration so + # saved checkpoints describe the geometry that is actually trained. + check_partial_rotary_factor(draft_config, normalize=True) draft_config.architectures = ["Qwen3DSparkModel"] draft_config.num_target_layers = num_target_layers draft_config.num_hidden_layers = num_draft_layers diff --git a/deepspec/modeling/dspark/qwen3/modeling.py b/deepspec/modeling/dspark/qwen3/modeling.py index 24d44416..342c98f1 100644 --- a/deepspec/modeling/dspark/qwen3/modeling.py +++ b/deepspec/modeling/dspark/qwen3/modeling.py @@ -21,6 +21,7 @@ AcceptRatePredictor, DSparkForwardOutput, build_eval_mask, + check_partial_rotary_factor, create_dspark_attention_mask, create_noise_embed, create_position_ids, @@ -236,6 +237,10 @@ def __init__(self, config) -> None: ] ) self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + # Warn on checkpoints whose config declares partial RoPE: this modeling + # rotates the full head dim, so the declaration is untrue and breaks + # serving engines that honor it (e.g. vLLM's dflash/dspark path). + check_partial_rotary_factor(config) self.rotary_emb = Qwen3RotaryEmbedding(config) self.fc = nn.Linear( len(self.target_layer_ids) * config.hidden_size,