feat(speculative): support Gemma-4-E4B as a streaming DFlash/DSpark target

modelopt/torch/speculative/plugins/modeling_fakebase.py

@@ -72,6 +72,40 @@
 _SAFETENSORS_SINGLE_FILENAMES = ["model.safetensors", "consolidated.safetensors"]
 
 
+def _resolve_rope_theta(base_cfg, attn_kind: str = "sliding_attention") -> float | None:
+    """Return the base model's RoPE theta, handling nested ``rope_parameters``.
+
+    Most models expose a flat ``rope_theta``. Gemma 4 instead nests per-attention-kind RoPE
+    settings under ``rope_parameters``, e.g.::
+
+        {"full_attention":    {"rope_theta": 1e6, "rope_type": "proportional",
+                               "partial_rotary_factor": 0.25},
+         "sliding_attention": {"rope_theta": 1e4, "rope_type": "default"}}
+
+    A flat ``getattr(base_cfg, "rope_theta", None)`` returns ``None`` there, and the draft then
+    silently trains on the draft class's default theta instead of the base's — training loss and
+    accuracy still improve while MT-Bench AAL is capped, because RoPE frequencies get baked into
+    the trained weights.
+
+    ``attn_kind`` selects which entry to read; it must match the attention the DRAFT uses. The
+    default is ``sliding_attention`` because SWA drafts are the common case for Gemma 4, and its
+    ``rope_type`` is plain ``default`` (the ``full_attention`` entry uses ``proportional`` rope
+    with ``partial_rotary_factor``, which the draft classes do not implement).
+    """
+    theta = getattr(base_cfg, "rope_theta", None)
+    if theta is not None:
+        return theta
+    params = getattr(base_cfg, "rope_parameters", None)
+    if not isinstance(params, dict):
+        return None
+    entry = params.get(attn_kind)
+    if entry is None:
+        # Single-kind nested form, or an unknown kind name: fall back to the sole entry.
+        values = [v for v in params.values() if isinstance(v, dict) and "rope_theta" in v]
+        entry = values[0] if len(values) == 1 else None
+    return entry.get("rope_theta") if isinstance(entry, dict) else None
+
+
 class FakeBaseConfig(PretrainedConfig):
     """Minimal config for FakeBaseModel that supports offline speculative decoding training."""
 
@@ -203,7 +237,7 @@ def from_source(cls, source: str, trust_remote_code: bool = False) -> "FakeBaseM
             num_key_value_heads=getattr(base_cfg, "num_key_value_heads", None),
             intermediate_size=getattr(base_cfg, "intermediate_size", None),
             rms_norm_eps=getattr(base_cfg, "rms_norm_eps", 1e-6),
-            rope_theta=getattr(base_cfg, "rope_theta", None),
+            rope_theta=_resolve_rope_theta(base_cfg),
             final_norm_type=_select_final_norm_type(
                 getattr(base_cfg, "model_type", None), base_cfg
             ),

modelopt/torch/speculative/plugins/modeling_final_norm.py

@@ -93,6 +93,14 @@ def extra_repr(self):
     # M3's final norm is always gemma-style; map it here too so a config that lost its
     # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1.
     "minimax_m3_vl_text": "gemma_rmsnorm",
+    # Gemma 4 VLM nests the LLM as text_config with model_type "gemma4_text"; from_source
+    # reads the NESTED config, so a "gemma4" key alone would never match. Verified numerically
+    # on gemma-4-E4B-it that Gemma4RMSNorm is plain ``normed * weight`` — NOT the ``(1 + weight)``
+    # form used by Gemma 2/3 — reproducing HF ``hidden_states[-1]`` at cos=0.999999 (vs 0.9719
+    # and maxabs_err 47.6 for the ``(1 + weight)`` form), so plain ``rmsnorm`` is correct here
+    # and ``gemma_rmsnorm`` would be wrong. Both keys listed so a text-only checkpoint works too.
+    "gemma4_text": "rmsnorm",
+    "gemma4": "rmsnorm",
     # gpt_oss intentionally DISABLED: GptOssRMSNorm uses an fp32 weight + multiply-then-cast,
     # unlike _FinalRMSNorm's bf16 weight, so reusing it would silently bias reconstructed logits.
     # Re-enable once a gpt_oss-style class (fp32 weight, multiply-then-cast) is in _FINAL_NORM_CLASSES.

modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml

@@ -0,0 +1,127 @@
+# DSpark (full-train, from scratch) recipe for Gemma-4-E4B-it.
+#
+# Streaming: the real Gemma-4-E4B-it base is served by vLLM; the trainer uses a
+# fake base (FakeBaseModel carries embed_tokens + the final norm). DSpark
+# reconstructs the base teacher distribution from the captured PRE-norm hidden
+# and re-applies the base final norm before lm_head.
+#
+# Gemma-4-E4B-specific notes (all verified on PDX 2026-08-12):
+#
+#  * FINAL NORM: Gemma 4 nests the LLM under text_config with model_type
+#    "gemma4_text", so modeling_final_norm.py needed BOTH "gemma4_text" and
+#    "gemma4" added to _FINAL_NORM_TYPE_BY_MODEL_TYPE. Verified numerically that
+#    Gemma4RMSNorm is plain `normed * weight` (NOT Gemma 2/3's `(1 + weight)`),
+#    reproducing HF hidden_states[-1] at cos=0.999999, so the existing
+#    _FinalRMSNorm is correct as-is.
+#
+#  * ROPE: Gemma 4 has NO flat `rope_theta`; it nests per-attention-kind settings
+#    under `rope_parameters` (full_attention: theta 1e6 + rope_type
+#    "proportional" + partial_rotary_factor 0.25; sliding_attention: theta 1e4 +
+#    rope_type "default"). modeling_fakebase.py needed _resolve_rope_theta() to
+#    read the nested form. hf_dflash.py ENFORCES rope_theta from the base config
+#    and overwrites any value set here, so this could NOT be fixed from the yaml.
+#    We take the sliding_attention entry -> rope_theta 10000.0, matching the SWA
+#    draft below. (The full_attention entry's "proportional" rope +
+#    partial_rotary_factor is not implemented in the draft classes.)
+#
+#  * CAPTURE IDS: vLLM's EagleModelMixin captures POST-layer with residual added
+#    and indexes as layer_idx+1, so valid ids are 1..42 and 42 is the TRUE final
+#    layer (verified: cos(id42, final_norm_INPUT) = 1.0000, and ids 9/18/27/36
+#    match HF hidden_states[id] at cos=1.0000 with off-by-one dropping to
+#    0.55-0.98). Gemma 4 repeats 5x sliding + 1x full attention, so the
+#    full_attention layers (0-based [5,11,17,23,29,35,41]) correspond to capture
+#    ids [6,12,18,24,30,36,42]; we sample those to land on residual-stream
+#    boundaries rather than spacing uniformly.
+#
+#  * MASK TOKEN: Gemma 4's vocab is fully packed (no free/unused ids), but it
+#    ships a native `<mask>` token at id 4 — used directly.
+
+metadata:
+  recipe_type: speculative_dflash
+  description: DSpark (DFlash backbone + Markov + confidence head) for Gemma-4-E4B-it, SWA draft.
+
+model:
+  model_name_or_path:
+  trust_remote_code: true
+  use_fake_base_for_offline: true
+
+data:
+  mode: streaming
+  data_path:
+  offline_data_path:
+  chat_template:
+
+training:
+  output_dir:
+  num_train_epochs: 1
+  per_device_train_batch_size: 4
+  gradient_accumulation_steps: 1
+  learning_rate: 1.0e-4
+  warmup_steps: 500
+  training_seq_len: 4096
+  logging_steps: 20
+  save_steps: 1000
+  cp_size: 1
+  dp_shard_size: 1
+  disable_tqdm: true
+  # Eval runs the DFlash backbone only (Markov head not applied in eval forward),
+  # so AR would misreport. Compare via export + offline AL harness instead.
+  estimate_ar: false
+  ar_validate_steps: 0
+  answer_only_loss: true
+  do_eval: false
+  lr_scheduler_type: linear
+  save_strategy: steps
+  weight_decay: 0.0
+  max_grad_norm: 1.0
+  dataloader_drop_last: true
+  bf16: true
+  tf32: true
+  remove_unused_columns: false
+  ddp_find_unused_parameters: true
+  ddp_timeout: 1800
+  report_to: none
+
+dflash:
+  dflash_block_size: 8
+  dflash_num_anchors: 512
+  dflash_use_torch_compile: false
+  dflash_self_logit_distillation: false
+  # block_size=8 -> decay gamma 4 (matches the K2.6 DSpark regime).
+  dflash_loss_decay_factor: 4.0
+  # Gemma 4 ships a native <mask> token at id 4 (vocab is fully packed, so there
+  # is no spare/unused id to borrow the way Kimi's 163838 was).
+  dflash_mask_token_id: 4
+  # --- DSpark three-term loss (DeepSpec L1/TVD-dominant defaults) ---
+  dflash_ce_loss_alpha: 0.1
+  dflash_l1_loss_alpha: 0.9
+  dflash_confidence_head_alpha: 1.0
+  dflash_architecture_config:
+    # Draft dims are set explicitly — the draft is an independent model and does
+    # NOT inherit these from the base (hidden_size/vocab/rope_theta ARE forced to
+    # the base and need not be set here).
+    num_hidden_layers: 5
+    num_attention_heads: 16
+    num_key_value_heads: 4
+    head_dim: 256
+    intermediate_size: 10240
+    projector_type: dspark
+    # Markov head: low-rank first-order transition bias, memoryless variant.
+    markov_rank: 256
+    markov_head_type: vanilla
+    use_confidence_head: true
+    # --- SWA draft (user decision 2026-08-12: try SWA first) ---
+    # DFlashAttention enables sliding-window attention only when the draft config
+    # carries BOTH `layer_types` and `sliding_window`; it then applies the window
+    # on layers whose layer_types entry is "sliding_attention". Matching the base's
+    # window of 512 and its sliding-layer rope_theta of 10000.
+    # NOTE: K3 measured SWA costing ~23% AL vs full attention at window 1024, and
+    # this window is smaller still — expect an AL hit and compare against a
+    # full-attention control before concluding.
+    sliding_window: 512
+    layer_types:
+      - sliding_attention
+      - sliding_attention
+      - sliding_attention
+      - sliding_attention
+      - sliding_attention