From 1d321b2f7cd2d853f129829b374bfc90568bcf9f Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Thu, 13 Aug 2026 12:19:29 +0000 Subject: [PATCH 1/8] feat(speculative): support Gemma-4-E4B as a streaming DFlash/DSpark target Adds the pieces needed to train a drafter against Gemma-4-E4B-it with streaming hidden-state capture, plus two fake-base fixes the model exposed. Validated end-to-end on AWS-PDX: 20-step DSpark streaming smoke, loss 3.79 -> 3.28 monotonically, drafter exported (62 tensors). modeling_final_norm: whitelist "gemma4_text" / "gemma4". Gemma 4 is a VLM that nests the LLM under text_config with model_type "gemma4_text", so from_source reads the nested config and a "gemma4" key alone would never match. Without an entry the fake base builds no final norm and the streaming teacher logits are reconstructed from an un-normed hidden -- a silent distillation-target corruption. Verified numerically that Gemma4RMSNorm is plain `normed * weight`, NOT the `(1 + weight)` form used by Gemma 2/3: it reproduces HF hidden_states[-1] at cos=0.999999, versus cos=0.9719 / maxabs_err 47.6 for the `(1 + weight)` form. So plain "rmsnorm" is correct here and "gemma_rmsnorm" would be wrong. modeling_fakebase: resolve RoPE theta from nested rope_parameters. Gemma 4 has no flat `rope_theta`; it nests per-attention-kind settings under `rope_parameters` (full_attention: 1e6 + rope_type "proportional" + partial_rotary_factor 0.25; sliding_attention: 1e4 + rope_type "default"). The flat getattr returned None, so the draft would silently train on its own class default -- loss and accuracy still improve while MT-Bench AAL is capped, because RoPE frequencies bake into the trained weights. This could not be worked around from the recipe: hf_dflash enforces rope_theta from the base config and overwrites any dflash_architecture_config value. Models with a flat rope_theta are unaffected (covered by the added fallbacks). chat_template_train.jinja: Gemma 4's stock template has no generation markers, so `return_assistant_tokens_mask` yields an all-zero loss_mask under answer_only_loss and EVERY row is rejected with "no fetchable sample found in the entire corpus". This copy wraps the model-turn content in generation markers; the rendered text is byte-identical to the stock template, and the zero-mask rate drops from 200/200 to 4/200 at max_seq_len 2048. dspark_gemma4_e4b.yaml / hf_streaming_dspark_smoke.yaml: DSpark recipe and a single-node co-located streaming smoke. Notable Gemma-4 settings are the native token id 4 (the vocab is fully packed, so there is no spare id to borrow), an SWA draft matching the base's 512 window, and capture ids [6,12,18,24,36,42] -- the full-attention layers of the 5:1 sliding/full cycle. len(EAGLE_CAPTURE_IDS) must equal num_draft_layers + 1, since the projector is sized from the draft's num_hidden_layers. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative/plugins/modeling_fakebase.py | 36 +- .../plugins/modeling_final_norm.py | 8 + .../dspark_gemma4_e4b.yaml | 127 ++++++ .../gemma-4-E4B-it/chat_template_train.jinja | 390 ++++++++++++++++++ .../hf_streaming_dspark_smoke.yaml | 98 +++++ 5 files changed, 658 insertions(+), 1 deletion(-) create mode 100644 modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml create mode 100644 tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja create mode 100644 tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml diff --git a/modelopt/torch/speculative/plugins/modeling_fakebase.py b/modelopt/torch/speculative/plugins/modeling_fakebase.py index 2b5fe989c03..ac787c789f1 100644 --- a/modelopt/torch/speculative/plugins/modeling_fakebase.py +++ b/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 ), diff --git a/modelopt/torch/speculative/plugins/modeling_final_norm.py b/modelopt/torch/speculative/plugins/modeling_final_norm.py index 718d591b662..9e4a7ad6dd0 100644 --- a/modelopt/torch/speculative/plugins/modeling_final_norm.py +++ b/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. diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml new file mode 100644 index 00000000000..981141f2d59 --- /dev/null +++ b/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 `` 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 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 diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja b/tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja new file mode 100644 index 00000000000..f4b962ede6c --- /dev/null +++ b/tools/launcher/examples/google/gemma-4-E4B-it/chat_template_train.jinja @@ -0,0 +1,390 @@ +{# + Template: Google Gemma 4 Canonical Chat Template + Author: Google Gemma Engineering Team + Published: 2026-07-09 + Context: Fixed tool-calling loops, turn closures, and thinking content-ordering. +#} +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%} + {%- if thinking_text and thinking_gate -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- else -%} + {{- raise_exception( + "chat_template: tool_calls[].function.arguments must be a " + "JSON object (mapping), not a string. Deserialize arguments " + "before passing to the template." + ) -}} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {%- if role == 'model' -%} + {%- generation -%}{{- captured_content -}}{%- endgeneration -%} + {%- else -%} + {{- captured_content -}} + {%- endif -%} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and (not message.get('tool_calls') or ns_tr_out.flag) + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} +{%- endif -%} diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml new file mode 100644 index 00000000000..5ff232145ea --- /dev/null +++ b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke.yaml @@ -0,0 +1,98 @@ +# DSpark streaming SMOKE run (co-located single node) for Gemma-4-E4B-it on AWS-PDX. +# +# Purpose: prove the streaming chain end-to-end (vLLM serve -> aux hidden-state +# capture -> NIXL transfer -> fake-base trainer) for a NEW base model. Only ~20 +# steps on a 1024-row corpus; the numbers are meaningless, the point is that the +# pipeline runs and the loss is finite and decreasing. +# +# Topology: serve TP=1 on GPU 0 (E4B is 16 GB, fits one B300), DSpark trainer on +# GPUs 4-7. Intra-node NIXL, no cross-node EFA. +# +# Gemma-4-E4B specifics verified on PDX 2026-08-12: +# * EAGLE_CAPTURE_IDS: vLLM's EagleModelMixin captures POST-layer with the +# residual added and indexes as layer_idx+1, so valid ids are 1..42 and 42 is +# the TRUE final layer (verified cos=1.0000 against HF hidden_states, with +# off-by-one dropping to 0.55-0.98). Gemma 4 repeats 5x sliding + 1x full +# attention, so full-attention layers sit at capture ids +# [6,12,18,24,30,36,42]; we sample those to land on residual-stream +# boundaries instead of spacing uniformly (deep layers are near-redundant: +# adjacent cosine ~0.98-0.99 around id 36 vs 0.55-0.70 around id 9). +# * NO --trust-remote-code needed (the repo ships no .py) and no vLLM patch: +# gemma4 is natively supported in the 2026-08-11 nightly container. +# * Corpus has its user-only `messages` column DROPPED — hf_streaming_dataset +# prefers `messages` over `conversations` and a user-only one makes streaming +# SILENTLY HANG. + +job_name: Gemma-4-E4B_DSpark_streaming_smoke +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/gemma-4-E4B-it + + task_0: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.mode=streaming + # Gemma 4's stock chat template has NO generation markers, so + # answer_only_loss silently yields an all-zero loss_mask and EVERY row is + # rejected ("no fetchable sample found in the entire corpus"). This copy + # adds them; see the template file header for details. + - data.chat_template=examples/google/gemma-4-E4B-it/chat_template_train.jinja + - data.data_path=/smokedata + - training.output_dir=/scratchspace/dspark_smoke + - training.training_seq_len=2048 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.max_steps=20 + - training.logging_steps=1 + - training.save_steps=20 + - training.per_device_train_batch_size=1 + - training.answer_only_loss=true + - training.report_to=none + environment: + - HF_MODEL_CKPT: <> + # MUST be exactly num_draft_layers + 1 entries (5 draft layers -> 6 ids): + # the projector is sized from the DRAFT's num_hidden_layers, so 7 ids gave + # "mat1 and mat2 shapes cannot be multiplied (2048x15360 and 12800x2560)". + # Confirmed against the working runs: K2.6 6 layers -> 7 ids, gpt-oss 5 -> 6. + # Chosen from the full-attention layers of the 5:1 sliding/full cycle + # ([6,12,18,24,30,36,42]), dropping 30 to keep both ends and the true final 42. + - EAGLE_CAPTURE_IDS: "[6,12,18,24,36,42]" + - SERVE_GPU: "0" + - SERVE_TP: "1" + - SERVE_GPU_MEM_UTIL: "0.85" + - STREAMING_NUM_WORKERS: "1" + - SERVE_MAX_MODEL_LEN: "2176" + - SERVE_MAX_NUM_SEQS: "4" + - SERVE_READY_TIMEOUT: "2400" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + # NIXL hidden-state transport. Without these, NIXL falls back to the UCX + # backend, which reports "8 NVIDIA GPU(s) were detected, but UCX CUDA + # support was not found" and then dies with NIXL_ERR_REMOTE_DISCONNECT the + # moment the trainer tries to pull hidden states. LIBFABRIC+efa is what the + # working Kimi-K2.6 streaming runs use. + - NIXL_BACKENDS: LIBFABRIC + - FI_PROVIDER: efa + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + # The 2026-08-11 nem35 image has NO libfabric and no NIXL plugins, so + # NIXL_BACKENDS=LIBFABRIC dies with NIXL_ERR_NOT_FOUND and the UCX fallback + # dies with NIXL_ERR_REMOTE_DISCONNECT ("UCX CUDA support was not found"). + # The auxfix image carries the AWS libfabric stack and is what the working + # Kimi-K2.6 streaming runs use. Its vLLM is older (transformers 5.12.1) -- + # gemma4 support must be re-verified in THIS image. + container: /home/haoguo/lustre/containers/vllm-nightly-efa-x86_64-auxfix.sqsh + container_mounts: + - /home/haoguo/lustre/hf-local:/hf-local + - /home/haoguo/lustre/g4_smoke_corpus:/smokedata From 7b56dd6e33b04a5eca4a6a4dbf87312d2144c750 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:40:53 +0000 Subject: [PATCH 2/8] feat(speculative): Gemma4-style DSpark draft block (sandwich norms, k_eq_v) Aligns the Gemma-4 draft with `deepseek-ai/dspark_gemma4_12b_block7`, the reference checkpoint vLLM's `gemma4_dspark.py` was written for (vLLM #47216). Without this the exported drafter cannot be served by that path at all. Two things were wrong before, both silent: * Block shape. `Gemma4DSparkDecoderLayer` inherits `Gemma4MTPDecoderLayer`, which looks up `pre_feedforward_layernorm`, `post_feedforward_layernorm` and `layer_scalar` BY NAME. The Qwen3-style `DFlashDecoderLayer` has none of them (single pre-norm per sub-block), and the DSpark weight loader only fills names it finds -- anything missing stays randomly initialized with no error. So a Qwen3-shaped draft loads "successfully" and then produces garbage. `DFlashGemma4DecoderLayer` reproduces Gemma4's order exactly: norm -> attn -> norm -> +residual -> norm -> mlp -> norm -> +residual, then `* layer_scalar` (a buffer, matching vLLM's `register_buffer`). * Attention. Gemma4 sizes attention PER LAYER: full-attention layers use `global_head_dim`, and under `attention_k_eq_v` also `num_global_key_value_heads` (1 in the reference, vs 8 for sliding). Under k_eq_v there is no `v_proj` at all -- V is derived from the K projection and passed through a weightless `v_norm`. `DFlashGemma4Attention` mirrors vLLM's `gemma4_layer_config` for this. Mis-sizing `k_proj` is silent for the same reason as above. The draft layer class is selected from `model_type`, so every non-Gemma4 draft keeps the existing Qwen3-style block unchanged. Recipe `dspark_gemma4_e4b.yaml` now matches the reference backbone: 5 full-attention layers, `attention_k_eq_v: true`, `global_head_dim: 512`, `num_global_key_value_heads: 1`. This replaces the earlier SWA draft, which vLLM cannot serve: `_build_fused_kv_buffers()` runs unconditionally at the end of `load_weights` and asserts every draft layer has `use_k_eq_v`, failing with "Gemma4 DSpark fused precompute assumes uniform attention_k_eq_v layers". Also adds `convert_gemma4_dspark_to_vllm.py`, which rewrites an export into the layout vLLM expects: `architectures` -> `Gemma4DSparkModel`, `model_type` -> `gemma4_text`, promotes `target_layer_ids` / `markov_rank` / `mask_token_id` from the nested `dflash_config` to the top level, renames `markov_w1/w2` under `markov_head.`, and bakes in `lm_head` + `embed_tokens` from the base (Gemma 4 is tie_word_embeddings, and the export ships neither -- another parameter the loader would otherwise leave random). Verified: the built backbone matches the reference checkpoint's safetensors header exactly -- all 16 backbone tensors identical in name and shape, no extra tensors, and no `v_proj` under k_eq_v. --- .../export/convert_gemma4_dspark_to_vllm.py | 103 +++++++++++ .../speculative/plugins/modeling_dflash.py | 170 +++++++++++++++++- .../dspark_gemma4_e4b.yaml | 42 +++-- 3 files changed, 295 insertions(+), 20 deletions(-) create mode 100644 examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py diff --git a/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py new file mode 100644 index 00000000000..ae8a8bfc203 --- /dev/null +++ b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py @@ -0,0 +1,103 @@ +#!/usr/bin/env python3 +"""Convert a ModelOpt DSpark drafter export into the layout vLLM's +Gemma4DSparkForCausalLM expects. + +Three classes of mismatch, all on the ModelOpt side (vLLM needs no patch): + +1. config.json + - architectures: DFlashDraftModel -> Gemma4DSparkModel (registry key) + - model_type: qwen3 -> gemma4_text (Gemma4DSparkAttention + reads layer_types/head_dim/global_head_dim off a Gemma4 text config) + - target_layer_ids / markov_rank are read as TOP-LEVEL attrs by + Gemma4DSparkModel.__init__, but ModelOpt nests them under dflash_config. + +2. weight names + - markov_w1/markov_w2 -> markov_head.markov_w1/markov_head.markov_w2 + (DSparkMarkovHead registers them under a `markov_head` submodule) + +3. missing tensors <-- the one that silently destroys AL + - Gemma4DSparkForCausalLM builds its OWN ParallelLMHead and VocabParallelEmbedding + and its load_weights() only fills names it finds; anything absent stays + RANDOMLY INITIALIZED with no error. The ModelOpt export ships neither + lm_head nor embed_tokens, so both must be baked in from the base model. + Gemma 4 is tie_word_embeddings=true, so both come from the same tensor: + model.language_model.embed_tokens.weight +""" + +import argparse +import json +import os +import shutil + +from safetensors.torch import load_file, save_file + +ap = argparse.ArgumentParser() +ap.add_argument("--drafter", required=True, help="ModelOpt export dir") +ap.add_argument("--base", required=True, help="base Gemma-4-E4B-it dir (for lm_head/embed)") +ap.add_argument("--out", required=True) +args = ap.parse_args() + +os.makedirs(args.out, exist_ok=True) +cfg = json.load(open(os.path.join(args.drafter, "config.json"))) +df = cfg.get("dflash_config", {}) or {} + +new = dict(cfg) +new["architectures"] = ["Gemma4DSparkModel"] +new["model_type"] = "gemma4_text" +# promote what Gemma4DSparkModel.__init__ reads off the top level +new["target_layer_ids"] = df.get("target_layer_ids", cfg.get("target_layer_ids")) +new["markov_rank"] = df.get("markov_rank", 256) +new["mask_token_id"] = df.get("mask_token_id") +new["block_size"] = cfg.get("block_size") +new.setdefault("draft_vocab_size", cfg["vocab_size"]) +# Gemma4 attention knobs consulted by gemma4_layer_config()/Gemma4DSparkAttention +new.setdefault("global_head_dim", cfg["head_dim"]) +new.setdefault("attention_k_eq_v", False) +new.setdefault("sliding_window", 512) +new.setdefault("final_logit_softcapping", None) +json.dump(new, open(os.path.join(args.out, "config.json"), "w"), indent=2) +print( + "config: architectures={} model_type={} target_layer_ids={} markov_rank={}".format( + new["architectures"], new["model_type"], new["target_layer_ids"], new["markov_rank"] + ) +) + +sd = load_file(os.path.join(args.drafter, "model.safetensors")) +print(f"loaded {len(sd)} drafter tensors") + +out = {} +renamed = 0 +for k, v in sd.items(): + nk = k + if k.startswith(("markov_w1", "markov_w2")): + nk = "markov_head." + k + renamed += 1 + out[nk] = v +print(f"renamed {renamed} markov tensors -> markov_head.*") + +# --- bake in lm_head + embed_tokens from the base (tied on Gemma 4) --- +idx_path = os.path.join(args.base, "model.safetensors.index.json") +if os.path.exists(idx_path): + wm = json.load(open(idx_path))["weight_map"] + key = next(k for k in wm if k.endswith("language_model.embed_tokens.weight")) + base_sd = load_file(os.path.join(args.base, wm[key])) +else: + base_sd = load_file(os.path.join(args.base, "model.safetensors")) + key = next(k for k in base_sd if k.endswith("language_model.embed_tokens.weight")) +emb = base_sd[key] +print(f"base embed_tokens {tuple(emb.shape)} from {key!r}") +assert emb.shape[0] == cfg["vocab_size"] and emb.shape[1] == cfg["hidden_size"], ( + f"base embed {tuple(emb.shape)} does not match draft vocab/hidden " + f"({cfg['vocab_size']},{cfg['hidden_size']})" +) + +out["lm_head.weight"] = emb.clone() +out["embed_tokens.weight"] = emb.clone() +print("baked lm_head.weight + embed_tokens.weight (tie_word_embeddings=true on Gemma 4)") + +save_file(out, os.path.join(args.out, "model.safetensors"), metadata={"format": "pt"}) +for f in ("tokenizer.json", "tokenizer_config.json"): + src = os.path.join(args.base, f) + if os.path.exists(src): + shutil.copy(src, os.path.join(args.out, f)) +print(f"wrote {len(out)} tensors -> {args.out}") diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 6463cb4109d..45f895c14c2 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -221,6 +221,117 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m return self.o_proj(attn_output) +class DFlashGemma4Attention(DFlashAttention): + """DFlash attention for a Gemma4-style draft. + + Two deltas versus the Qwen3-style :class:`DFlashAttention`: + + * ``attention_k_eq_v``: Gemma4 can derive V from the K projection instead of + carrying a separate ``v_proj``, halving the KV parameters. vLLM's + ``Gemma4DSparkAttention`` does exactly this (``v_src = k`` when + ``use_k_eq_v``), and its fused context-KV precompute *asserts* every draft + layer is built this way, so a draft trained with a separate ``v_proj`` + cannot be served by that path at all. + * ``v_norm``: applied to V, with **no learnable weight**, mirroring vLLM's + ``RMSNorm(..., has_weight=False)``. Plain (Qwen3) DFlash does not norm V. + + ``use_k_eq_v`` follows vLLM: full-attention layers only, and only when the + config opts in. Sliding layers keep their own ``v_proj``. + """ + + def __init__(self, config, layer_idx): + """Initialize Gemma4 draft attention with per-layer dims, dropping ``v_proj`` under k_eq_v.""" + super().__init__(config, layer_idx) + layer_types = getattr(config, "layer_types", None) + is_full = layer_types is None or layer_types[layer_idx] == "full_attention" + self.use_k_eq_v = is_full and getattr(config, "attention_k_eq_v", False) + + # Gemma4's attention dims are PER LAYER: full-attention layers use a larger + # ``global_head_dim`` and, under k_eq_v, a smaller ``num_global_key_value_heads``. + # This mirrors vLLM's ``gemma4_layer_config`` (transformers_utils/configs/gemma4.py), + # which Gemma4DSparkAttention calls to size q/k/o. Getting this wrong is silent: + # the DSpark weight loader only fills names it finds, so a mis-shaped k_proj is + # simply left randomly initialized. + if is_full: + self.head_dim = getattr(config, "global_head_dim", None) or self.head_dim + if self.use_k_eq_v: + self.num_kv_heads = ( + getattr(config, "num_global_key_value_heads", None) or self.num_kv_heads + ) + self.num_key_value_groups = self.num_heads // self.num_kv_heads + self.scaling = self.head_dim**-0.5 + attn_bias = getattr(config, "attention_bias", False) + self.q_proj = nn.Linear( + config.hidden_size, self.num_heads * self.head_dim, bias=attn_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, self.num_kv_heads * self.head_dim, bias=attn_bias + ) + self.o_proj = nn.Linear( + self.num_heads * self.head_dim, config.hidden_size, bias=attn_bias + ) + self.q_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps) + + if self.use_k_eq_v: + # Registered by the parent; drop it so it is neither trained nor exported. + del self.v_proj + self.v_proj = None + elif is_full: + self.v_proj = nn.Linear( + config.hidden_size, + self.num_kv_heads * self.head_dim, + bias=getattr(config, "attention_bias", False), + ) + # vLLM builds this as ``RMSNorm(..., has_weight=False)`` and the reference + # checkpoint ships NO v_norm tensor, so keep the scale fixed at ones and + # non-persistent: it must not appear in the exported state_dict. + self.v_norm = _NORM_CLS(self.head_dim, eps=config.rms_norm_eps) + del self.v_norm.weight + self.v_norm.register_buffer("weight", torch.ones(self.head_dim), persistent=False) + + def _project_v(self, target_hidden, hidden_states, k_ctx, k_noise): + """Return the V sequence, from K under k_eq_v or from ``v_proj`` otherwise.""" + if self.use_k_eq_v: + return k_ctx, k_noise + return self.v_proj(target_hidden), self.v_proj(hidden_states) + + def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None): + """Forward with KV injection; V is normed and, under k_eq_v, shares K's projection.""" + bsz, q_len, _ = hidden_states.shape + ctx_len = target_hidden.shape[1] + + q = self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim) + q = self.q_norm(q).transpose(1, 2) + + k_ctx = self.k_proj(target_hidden) + k_noise = self.k_proj(hidden_states) + k = torch.cat([k_ctx, k_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + k = self.k_norm(k).transpose(1, 2) + + v_ctx, v_noise = self._project_v(target_hidden, hidden_states, k_ctx, k_noise) + v = torch.cat([v_ctx, v_noise], dim=1).view(bsz, ctx_len + q_len, -1, self.head_dim) + # vLLM norms V (no RoPE on V), unlike the Qwen3-style path. + v = self.v_norm(v).transpose(1, 2) + + cos, sin = position_embeddings + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + attn_fn = self._get_attn_fn() + attn_output, _ = attn_fn( + self, + q, + k, + v, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, + ) + attn_output = attn_output.reshape(bsz, q_len, -1) + return self.o_proj(attn_output) + + class DFlashDecoderLayer(nn.Module): """Draft decoder layer with KV injection.""" @@ -248,6 +359,56 @@ def forward(self, hidden_states, target_hidden, position_embeddings, attention_m return hidden_states +class DFlashGemma4DecoderLayer(nn.Module): + """Draft decoder layer matching Gemma4's block, with KV injection. + + Gemma4 wraps each sub-block in a *pair* of norms ("sandwich norm") and scales + the layer output by a learned ``layer_scalar``, where Qwen3 uses a single + pre-norm per sub-block. vLLM's ``Gemma4MTPDecoderLayer`` -- which + ``Gemma4DSparkDecoderLayer`` inherits -- looks up + ``pre_feedforward_layernorm`` / ``post_feedforward_layernorm`` / + ``layer_scalar`` by name, and its DSpark weight loader silently leaves any + parameter it cannot find randomly initialized. A Qwen3-shaped draft + therefore *loads without error* and produces garbage, so the shapes must + match exactly. + + The residual/norm order below mirrors ``Gemma4MTPDecoderLayer.forward``: + norm -> attn -> norm -> +residual -> norm -> mlp -> norm -> +residual, then + scale by ``layer_scalar``. + """ + + def __init__(self, config, layer_idx): + """Initialize a Gemma4-style draft layer (sandwich norms + layer scalar).""" + super().__init__() + self.self_attn = DFlashGemma4Attention(config, layer_idx) + self.mlp = _MLP_CLS(config) + self.input_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) + self.pre_feedforward_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) + self.post_feedforward_layernorm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) + # A buffer (not a parameter) to match vLLM's `register_buffer`, so the + # exported tensor name and shape line up with the reference checkpoint. + self.register_buffer("layer_scalar", torch.ones(1)) + + def forward(self, hidden_states, target_hidden, position_embeddings, attention_mask=None): + """Forward with sandwich norms, KV injection, and the layer scalar.""" + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states, target_hidden, position_embeddings, attention_mask + ) + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = hidden_states + residual + + residual = hidden_states + hidden_states = self.pre_feedforward_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = self.post_feedforward_layernorm(hidden_states) + hidden_states = hidden_states + residual + + return hidden_states * self.layer_scalar + + class DFlashModule(nn.Module): """DFlash draft module using Qwen3 components (MLP, RMSNorm, RotaryEmbedding).""" @@ -263,8 +424,15 @@ def __init__(self, config): self.hidden_norm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) # Decoder layers + # Gemma4 drafts need Gemma4's block shape (sandwich norms + layer_scalar, + # optional k_eq_v); everything else keeps the Qwen3-style block. + layer_cls = ( + DFlashGemma4DecoderLayer + if str(getattr(config, "model_type", "")).startswith("gemma4") + else DFlashDecoderLayer + ) self.layers = nn.ModuleList( - [DFlashDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + [layer_cls(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] ) self.norm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) self._rotary_config = config # Used by _maybe_init_rotary_emb diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml index 981141f2d59..bcaa443e201 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml @@ -97,31 +97,35 @@ dflash: 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). + # Aligned with the official reference drafter deepseek-ai/dspark_gemma4_12b_block7 + # (the checkpoint vLLM's gemma4_dspark.py was written for). Its backbone is + # 5 full-attention layers with attention_k_eq_v; vLLM's fused context-KV + # precompute ASSERTS every draft layer is built that way, so an SWA draft + # cannot be served by that path at all. num_hidden_layers: 5 num_attention_heads: 16 - num_key_value_heads: 4 + num_key_value_heads: 8 + # Full-attention layers use global_head_dim (512) and, under k_eq_v, + # num_global_key_value_heads (1) -- see gemma4_layer_config in vLLM. head_dim: 256 + global_head_dim: 512 + num_global_key_value_heads: 1 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 + # --- Gemma4 block shape (sandwich norms + layer_scalar) + k_eq_v --- + # model_type gemma4* selects DFlashGemma4DecoderLayer, which adds + # pre/post_feedforward_layernorm and layer_scalar. vLLM's + # Gemma4MTPDecoderLayer looks those up BY NAME and its DSpark loader leaves + # anything it cannot find randomly initialized -- silently. A Qwen3-shaped + # draft therefore loads without error and produces garbage. + model_type: gemma4_text + attention_k_eq_v: true layer_types: - - sliding_attention - - sliding_attention - - sliding_attention - - sliding_attention - - sliding_attention + - full_attention + - full_attention + - full_attention + - full_attention + - full_attention From 1dfab24859b68ca45a27bfe8d699265355a25089 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:56:04 +0000 Subject: [PATCH 3/8] feat(speculative): Gemma4-shaped DFlash/DSpark draft layer The Gemma-4 DSpark recipe added earlier produced a Qwen3-shaped draft, which vLLM's Gemma4 DSpark path cannot serve. Add the Gemma4 block shape so the exported drafter matches deepseek-ai/dspark_gemma4_12b_block7, the reference checkpoint that path was written for (vllm PR #47216). Every mismatch here fails SILENTLY: Gemma4DSparkForCausalLM.load_weights only fills parameter names it finds and leaves the rest randomly initialized, so a mis-shaped draft loads without error and produces garbage. DFlashGemma4DecoderLayer -- Gemma4 wraps each sub-block in a pair of norms ("sandwich norm") and scales the layer output by a learned layer_scalar, where Qwen3 uses a single pre-norm per sub-block. vLLM's Gemma4MTPDecoderLayer, which Gemma4DSparkDecoderLayer inherits, looks up pre_feedforward_layernorm / post_feedforward_layernorm / layer_scalar by name. The residual and norm order mirrors that forward exactly. DFlashGemma4Attention -- two deltas versus the Qwen3-style attention: * attention_k_eq_v derives V from the K projection instead of carrying a separate v_proj. vLLM's fused context-KV precompute asserts every draft layer is built this way, so a draft with its own v_proj cannot be served at all. * Gemma4 attention is heterogeneous: full-attention layers use global_head_dim and, under k_eq_v, num_global_key_value_heads. This mirrors vLLM's gemma4_layer_config(); without it k_proj comes out 8x too large (the reference drafter has k_proj [512, 3840] = 1 head x 512, not 8 x 512). V is normed with a fixed, non-persistent scale, matching vLLM's RMSNorm(has_weight=False); the reference checkpoint ships no v_norm tensor. Per-attention-kind RoPE -- full-attention and sliding layers use different head dims AND different rope_parameters (theta 1e6 vs 1e4), so one shared rotary module silently mismatches the head dim on one of the two kinds. vLLM builds RoPE per layer for this reason; DFlashModule now builds one per distinct layer_types entry and dispatches by layer. Selection is keyed on model_type starting with "gemma4", so every other model keeps the Qwen3-style block unchanged. Verified against the reference checkpoint's safetensors header: the draft backbone reproduces its tensor names and shapes exactly -- no missing, no extra, k_proj [512, 3840], layer_scalar [1], and no v_proj. Forward passes are finite for gemma4 full-attn with and without k_eq_v, gemma4 sliding, and qwen3. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../speculative/plugins/modeling_dflash.py | 53 +++++++++++++++++-- 1 file changed, 50 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 45f895c14c2..150be57f3a8 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -41,6 +41,7 @@ The draft architecture is independent of the target model. """ +import copy from dataclasses import dataclass import torch @@ -436,6 +437,8 @@ def __init__(self, config): ) self.norm = _NORM_CLS(config.hidden_size, eps=config.rms_norm_eps) self._rotary_config = config # Used by _maybe_init_rotary_emb + self._gemma4_rope_kinds = self._build_gemma4_rope_kinds(config) + self._layer_types = list(getattr(config, "layer_types", []) or []) # Explicit weight init is needed because DFlashModule is instantiated via # mtsp.convert() AFTER the base model's post_init() has already run, so HF's @@ -448,9 +451,46 @@ def _maybe_init_rotary_emb(self, device=None): Same pattern as EAGLE3's _maybe_init_rope. Avoids creating rotary_emb during __init__ (which runs on meta device during from_pretrained), preventing the meta-tensor inv_freq issue on checkpoint resume. + + Gemma4 needs one module PER attention kind, not one for the whole draft: + its full-attention layers use ``global_head_dim`` while sliding layers use + ``head_dim``, and the two kinds carry different ``rope_parameters`` (theta + 1e6 vs 1e4). vLLM builds RoPE per layer for exactly this reason; a single + shared module silently mismatches the head dim on one of the two kinds. """ if not hasattr(self, "rotary_emb"): self.rotary_emb = _ROTARY_CLS(config=self._rotary_config, device=device) + if self._gemma4_rope_kinds and not hasattr(self, "rotary_emb_by_kind"): + self.rotary_emb_by_kind = nn.ModuleDict( + { + kind: _ROTARY_CLS(config=cfg, device=device) + for kind, cfg in self._gemma4_rope_kinds.items() + } + ) + + @staticmethod + def _build_gemma4_rope_kinds(config): + """Per-attention-kind rotary configs for a Gemma4 draft, or ``{}`` otherwise. + + Returns a shallow copy of ``config`` per distinct ``layer_types`` entry with + ``head_dim`` and ``rope_parameters`` resolved for that kind. + """ + if not str(getattr(config, "model_type", "")).startswith("gemma4"): + return {} + layer_types = getattr(config, "layer_types", None) + if not layer_types: + return {} + rope_params = getattr(config, "rope_parameters", None) + kinds = {} + for kind in dict.fromkeys(layer_types): + cfg = copy.copy(config) + if kind == "full_attention": + cfg.head_dim = getattr(config, "global_head_dim", None) or config.head_dim + if isinstance(rope_params, dict) and isinstance(rope_params.get(kind), dict): + cfg.rope_parameters = dict(rope_params[kind]) + cfg.rope_theta = cfg.rope_parameters.get("rope_theta", config.rope_theta) + kinds[kind] = cfg + return kinds def _init_weights(self, config): """Initialize weights matching HF PreTrainedModel._init_weights.""" @@ -467,8 +507,15 @@ def forward(self, noise_embedding, target_hidden, position_ids, attention_mask=N target_hidden = self.hidden_norm(self.fc(target_hidden)) self._maybe_init_rotary_emb(device=hidden_states.device) position_embeddings = self.rotary_emb(hidden_states, position_ids) - - for layer in self.layers: - hidden_states = layer(hidden_states, target_hidden, position_embeddings, attention_mask) + per_kind = { + kind: emb(hidden_states, position_ids) + for kind, emb in getattr(self, "rotary_emb_by_kind", {}).items() + } + + for layer_idx, layer in enumerate(self.layers): + layer_pos = position_embeddings + if per_kind and layer_idx < len(self._layer_types): + layer_pos = per_kind.get(self._layer_types[layer_idx], position_embeddings) + hidden_states = layer(hidden_states, target_hidden, layer_pos, attention_mask) return self.norm(hidden_states) From b3dd61123e449542d4bb88cc800ebe25e43636f3 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:14:17 +0000 Subject: [PATCH 4/8] fix(speculative): export Gemma4 per-layer attention fields; harden the converter End-to-end validated on AWS-PDX: retrained the draft with the Gemma4 block, converted it, and served it under vLLM with the base Gemma-4-E4B-it. The engine comes up with the drafter attached and speculative decoding runs to completion. Two real gaps surfaced on the way. 1. The exporter dropped the fields that describe the draft's attention shapes. Gemma4 sizes attention PER LAYER: full-attention layers use `global_head_dim` and, under `attention_k_eq_v`, `num_global_key_value_heads` (see `gemma4_layer_config` in vLLM). The exported config carried only `head_dim` and `num_key_value_heads`, so vLLM rebuilt the draft with the SLIDING-layer dims and died with AssertionError: Attempted to load weight (torch.Size([512])) into parameter (torch.Size([256])) even though the trained weights were correct. The three fields are now propagated when present, so non-Gemma4 drafts are unaffected. 2. The converter hid that bug. It defaulted `global_head_dim` to `head_dim` and `attention_k_eq_v` to False, which silently produced a config describing a draft that does not exist. It now fails loudly if either is missing, or if k_eq_v is set without `num_global_key_value_heads`. The converter also now prints the serve command, because the draft attention backend is not discoverable: the draft re-runs backend auto-selection and lands on FLASH_ATTN, whose FA2 kernel caps head dimension at 256, while Gemma4 full-attention layers use 512 -- so serving fails with RuntimeError: FlashAttention forward only supports head dimension at most 256 The base model auto-selects FLASHINFER for exactly the same reason, but that choice is not inherited, and VLLM_ATTENTION_BACKEND does not reach the draft. It must be set via `speculative_config.attention_backend` (see vllm/v1/worker/gpu/spec_decode/dspark/utils.py). --- .../export/convert_gemma4_dspark_to_vllm.py | 29 +++++++++++++++++-- .../torch/export/plugins/hf_spec_export.py | 10 +++++++ 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py index ae8a8bfc203..fab427adcf0 100644 --- a/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py +++ b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py @@ -51,8 +51,20 @@ new["block_size"] = cfg.get("block_size") new.setdefault("draft_vocab_size", cfg["vocab_size"]) # Gemma4 attention knobs consulted by gemma4_layer_config()/Gemma4DSparkAttention -new.setdefault("global_head_dim", cfg["head_dim"]) -new.setdefault("attention_k_eq_v", False) +# Do NOT default these: Gemma4 sizes attention per layer, and quietly falling back to +# the sliding-layer values rebuilds the draft with the wrong q/k/o and q/k-norm shapes. +# They must come from the training config (see hf_spec_export.py). +for _req in ("global_head_dim", "attention_k_eq_v"): + if _req not in cfg: + raise SystemExit( + f"drafter config.json is missing {_req!r}. Re-export with a ModelOpt that " + "propagates the Gemma4 per-layer attention fields, or add it by hand." + ) +if cfg.get("attention_k_eq_v") and "num_global_key_value_heads" not in cfg: + raise SystemExit( + "attention_k_eq_v is set but num_global_key_value_heads is missing; " + "vLLM would size k_proj with the sliding-layer KV head count." + ) new.setdefault("sliding_window", 512) new.setdefault("final_logit_softcapping", None) json.dump(new, open(os.path.join(args.out, "config.json"), "w"), indent=2) @@ -101,3 +113,16 @@ if os.path.exists(src): shutil.copy(src, os.path.join(args.out, f)) print(f"wrote {len(out)} tensors -> {args.out}") +print() +print("Serve with (note the draft attention backend):") +_spec = ( + f'{{"model": "{args.out}", "num_speculative_tokens": 3, ' + '"method": "dspark", "attention_backend": "FLASHINFER"}' +) +print(f" vllm serve --speculative-config '{_spec}'") +print( + " FLASHINFER is REQUIRED: the draft re-runs backend auto-selection and lands on\n" + " FLASH_ATTN, whose FA2 kernel caps head dimension at 256, but Gemma4 full-attention\n" + " layers use global_head_dim=512. The VLLM_ATTENTION_BACKEND env var does NOT reach\n" + " the draft -- it is read from speculative_config.attention_backend." +) diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 255b1d9ab04..7d5f743bd7c 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -412,6 +412,16 @@ def _export_config(self): else: config["layer_types"] = ["full_attention"] * draft_config.num_hidden_layers + # Gemma4 sizes attention PER LAYER, so the serving side cannot reconstruct the + # draft's shapes from head_dim / num_key_value_heads alone: full-attention layers + # use ``global_head_dim`` and, under ``attention_k_eq_v``, ``num_global_key_value_heads`` + # (see gemma4_layer_config in vLLM). Omitting these makes vLLM rebuild the draft with + # the sliding-layer dims and fail with a shape mismatch on q/k/o and the q/k norms. + for _attr in ("global_head_dim", "num_global_key_value_heads", "attention_k_eq_v"): + _val = getattr(draft_config, _attr, None) + if _val is not None: + config[_attr] = _val + # Sliding-window attention: all draft layers use non-causal SWA (MiMo-style). vLLM's # _resolve_layer_attention reads dflash_config.use_swa + swa_window_size; with # layer_types left all "full_attention" it applies a non-causal sliding window to From e9448bb124a4a0fc78577f02056ea8c63eb1d7e1 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:45:20 +0000 Subject: [PATCH 5/8] docs(speculative): pin down Gemma4 DSpark draft attention semantics The Gemma-4-E4B DSpark path that trains, exports and serves end-to-end is a 5-layer all-full_attention, non-causal (bidirectional) draft with attention_k_eq_v, aligned with deepseek-ai/dspark_gemma4_12b_block7. Record the properties that were only implicit before: * Causality. vLLM resolves it per layer in _dflash_layer_causal(): an explicit dflash_config.causal overrides everything, else a layer is causal only when layer_types[i] == sliding_attention. This draft sets neither and is all full_attention, so all 5 layers are non-causal -- matching ModelOpt, whose DFlashAttention is non-causal and whose exporter emits no causal field. * FLASHINFER is required for two independent reasons, not one: global_head_dim 512 exceeds the FA2 head-dim cap of 256, AND the draft is non-causal so load_dspark_model sets use_non_causal=True. Either alone is fatal. * The Gemma4 SWA blocker is really a k_eq_v blocker. _build_fused_kv_buffers asserts uniform use_k_eq_v, and use_k_eq_v is only set on full_attention layers, so any sliding layer trips it. Mixed sliding/full would additionally require the V2 model runner. Also drop the stale SWA wording from the recipe header/description and explain why the local rope_theta is the right choice for an all-full draft. No functional change. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../export/convert_gemma4_dspark_to_vllm.py | 16 +++++-- .../speculative/plugins/modeling_dflash.py | 12 +++++ .../dspark_gemma4_e4b.yaml | 44 ++++++++++++++++--- 3 files changed, 61 insertions(+), 11 deletions(-) diff --git a/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py index fab427adcf0..968f35c2c3c 100644 --- a/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py +++ b/examples/speculative_decoding/export/convert_gemma4_dspark_to_vllm.py @@ -121,8 +121,16 @@ ) print(f" vllm serve --speculative-config '{_spec}'") print( - " FLASHINFER is REQUIRED: the draft re-runs backend auto-selection and lands on\n" - " FLASH_ATTN, whose FA2 kernel caps head dimension at 256, but Gemma4 full-attention\n" - " layers use global_head_dim=512. The VLLM_ATTENTION_BACKEND env var does NOT reach\n" - " the draft -- it is read from speculative_config.attention_backend." + " FLASHINFER is REQUIRED for TWO independent reasons; either alone is fatal:\n" + " 1. head dim -- the draft re-runs backend auto-selection and lands on FLASH_ATTN,\n" + " whose FA2 kernel caps head dimension at 256, but Gemma4 full-attention layers\n" + " use global_head_dim=512.\n" + " 2. causality -- the draft is NON-CAUSAL (bidirectional) on every layer, as DSpark\n" + " heads are: _dflash_layer_causal() only marks sliding_attention layers causal,\n" + " and this draft is all full_attention with no dflash_config.causal override.\n" + " load_dspark_model therefore sets use_non_causal=True, which the backend must\n" + " support. Runtime confirms with: 'Using FlashInfer for draft model non-causal\n" + " attention'.\n" + " The VLLM_ATTENTION_BACKEND env var does NOT reach the draft -- it is read from\n" + " speculative_config.attention_backend." ) diff --git a/modelopt/torch/speculative/plugins/modeling_dflash.py b/modelopt/torch/speculative/plugins/modeling_dflash.py index 150be57f3a8..8be96f2e11c 100644 --- a/modelopt/torch/speculative/plugins/modeling_dflash.py +++ b/modelopt/torch/speculative/plugins/modeling_dflash.py @@ -141,6 +141,18 @@ def __init__(self, config, layer_idx): self.num_key_value_groups = self.num_heads // self.num_kv_heads self.scaling = self.head_dim**-0.5 self.attention_dropout = getattr(config, "attention_dropout", 0.0) + # DFlash/DSpark drafts attend bidirectionally: a block of draft tokens is + # predicted in one shot, so those tokens must see each other. Serving must + # agree -- vLLM resolves per-layer causality in + # qwen3_dflash._dflash_layer_causal(): an explicit ``dflash_config.causal`` + # overrides all layers, otherwise a layer is causal only when + # ``layer_types[i] == "sliding_attention"``. The exporter emits no + # ``causal`` field for a plain full-attention draft, so it stays non-causal + # on both sides. With ``dflash_swa_window_size`` set, the exporter instead + # emits ``use_swa: True`` + an explicit ``causal: False`` and leaves + # ``layer_types`` all-full, which keeps vLLM non-causal too. Only mark + # layers ``sliding_attention`` if you also intend them to be CAUSAL at + # serving time, and train them that way -- see _build_draft_attention_mask. self.is_causal = False attn_bias = getattr(config, "attention_bias", False) diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml index bcaa443e201..6bbe32a9935 100644 --- a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml +++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml @@ -20,9 +20,12 @@ # 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.) +# We take the sliding_attention entry -> rope_theta 10000.0. Note this is the +# LOCAL theta even though the draft below is all full_attention: the +# full_attention entry's "proportional" rope + partial_rotary_factor 0.25 is +# not implemented in the draft classes, and the reference drafter +# (deepseek-ai/dspark_gemma4_12b_block7) also carries the default/1e4 form. +# The draft is trained and served with the same value, so it is self-consistent. # # * 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 @@ -35,10 +38,32 @@ # # * MASK TOKEN: Gemma 4's vocab is fully packed (no free/unused ids), but it # ships a native `` token at id 4 — used directly. +# +# * DRAFT ATTENTION IS NON-CAUSAL (bidirectional), like every other DSpark head. +# vLLM resolves this per layer in qwen3_dflash._dflash_layer_causal: an +# explicit `dflash_config.causal` overrides everything, otherwise a layer is +# causal only when layer_types[i] == "sliding_attention". This draft sets +# neither, and every layer is full_attention, so all 5 layers come out +# causal=False. That is what a block-parallel (semi-autoregressive) drafter +# wants: the block's tokens are predicted at once and must see each other. +# ModelOpt matches this -- DFlashAttention is non-causal, and the exporter +# does not emit a `causal` field. See "SERVING" below for the consequence. +# +# * SERVING requires speculative_config.attention_backend = "FLASHINFER", for +# TWO independent reasons; either alone is fatal: +# 1. head dim -- full_attention layers use global_head_dim=512 and the +# auto-selected FLASH_ATTN (FA2) caps head dimension at 256. +# 2. causality -- load_dspark_model sets attention_config.use_non_causal +# from dflash_has_any_non_causal(), which is True here, and the backend +# must be able to serve a non-causal mask. +# VLLM_ATTENTION_BACKEND does NOT reach the draft; it is read from +# speculative_config.attention_backend. Confirmed at runtime by the log line +# "Using FlashInfer for draft model non-causal attention". metadata: recipe_type: speculative_dflash - description: DSpark (DFlash backbone + Markov + confidence head) for Gemma-4-E4B-it, SWA draft. + description: DSpark (DFlash backbone + Markov + confidence head) for Gemma-4-E4B-it, + 5-layer full-attention draft aligned with the official reference checkpoint. model: model_name_or_path: @@ -99,9 +124,14 @@ dflash: dflash_architecture_config: # Aligned with the official reference drafter deepseek-ai/dspark_gemma4_12b_block7 # (the checkpoint vLLM's gemma4_dspark.py was written for). Its backbone is - # 5 full-attention layers with attention_k_eq_v; vLLM's fused context-KV - # precompute ASSERTS every draft layer is built that way, so an SWA draft - # cannot be served by that path at all. + # 5 full-attention layers with attention_k_eq_v. + # + # Why not SWA here: Gemma4DSparkModel._build_fused_kv_buffers asserts + # `all(a.use_k_eq_v for a in layers_attn)`, and Gemma4DSparkAttention only + # sets use_k_eq_v when layer_type == "full_attention". So the assert is + # really about k_eq_v uniformity, but because the two are coupled, ANY + # sliding layer trips it. Mixed sliding/full would additionally need the V2 + # model runner (see _resolve_layer_attention in qwen3_dflash.py). num_hidden_layers: 5 num_attention_heads: 16 num_key_value_heads: 8 From 5ccbef4f95686b23580f7fbff0b9d748d8ccc03a Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:50:52 +0000 Subject: [PATCH 6/8] feat(speculative): Gemma4 DSpark sliding-window (SWA) draft variant Adds the SWA variant of the Gemma-4-E4B DSpark recipe, plus the vLLM-side patch it needs, and documents why the two are separate. Training and export already support this with no code change: _build_draft_attention_mask() windows the CONTEXT (kv > q_real_pos - window) while leaving block-internal attention bidirectional, hf_dspark.py forwards dflash_swa_window_size, and hf_spec_export.py emits dflash_config.use_swa + swa_window_size + a top-level sliding_window with causal pinned False. So the recipe is a one-line delta over the validated full-attention baseline. Serving is the gap, and it is Gemma4-specific. Qwen3 DFlash layers resolve their window via _resolve_layer_attention(), which reads those fields. Gemma4DSparkAttention instead inherits Gemma4MTPAttention.__init__, which derives the window from the BASE layer pattern (layer_type == sliding_attention). Our draft is deliberately all full_attention -- that is what keeps attention_k_eq_v uniform for the fused-KV precompute assert -- so per_layer_sliding_window resolves to None on every layer, and an SWA-trained draft gets served with FULL attention. Nothing errors; acceptance just drops. gemma4_dspark_swa_vllm.patch fixes that by resolving the window through _resolve_layer_attention() and rebuilding self.attn when a window applies. Layers without a window are untouched, so the existing full-attention path is bit-identical. Verified by config simulation against the exporter output: all 5 layers get sliding_window=512, stay causal=False, keep use_k_eq_v=True (fused-KV assert passes), and layer_types stays uniform so the V2 model runner is not needed. The patch applies cleanly to gemma4_dspark.py. NOT yet run on GPU -- no acceptance-length number exists yet, and the recipe header says so. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../export/gemma4_dspark_swa_vllm.patch | 108 ++++++++++ .../dspark_gemma4_e4b_swa.yaml | 202 ++++++++++++++++++ 2 files changed, 310 insertions(+) create mode 100644 examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch create mode 100644 modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml diff --git a/examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch b/examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch new file mode 100644 index 00000000000..4f71eb7b5ca --- /dev/null +++ b/examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch @@ -0,0 +1,108 @@ +Make Gemma4 DSpark drafts honour dflash_config.use_swa / swa_window_size. + +Applies to (vLLM): vllm/model_executor/models/gemma4_dspark.py + +WHY + A DFlash/DSpark draft declares its sliding window through + dflash_config.use_swa + swa_window_size, and keeps layer_types all + "full_attention" on purpose -- that is what keeps attention_k_eq_v uniform, + which Gemma4DSparkModel._build_fused_kv_buffers() asserts. + + vLLM's Qwen3 DFlash layer resolves its window through + _resolve_layer_attention(), which reads exactly those fields. + Gemma4DSparkAttention does not: it inherits Gemma4MTPAttention.__init__, + which derives the window from the BASE model's layer pattern -- + + self.is_sliding = layer_type == "sliding_attention" + sliding_window = config.sliding_window if self.is_sliding else None + + so for an all-full_attention draft, per_layer_sliding_window is None on + every layer. A draft TRAINED with a window is then SERVED with full + attention. Nothing errors; acceptance length just quietly drops. + +WHAT + Resolve the window via _resolve_layer_attention(config, layer_idx) -- the + same call the Qwen3 DFlash layer makes -- and rebuild self.attn with + per_layer_sliding_window when a window applies. self.attn is already + replaced further down for the k_proj/v_proj rework, so rebuilding it here + follows the class's existing pattern; the stale static_forward_context entry + is popped first, exactly as Gemma4DSparkDecoderLayer does for + "{prefix}.self_attn.attn". + + Layers with no window are untouched, so the existing full-attention path + (deepseek-ai/dspark_gemma4_12b_block7 and the recipe it matches) is + bit-identical. + +VERIFIED (by config simulation, against the config ModelOpt's exporter emits +for modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml) + - all 5 layers get sliding_window=512 + - all 5 layers stay causal=False (dflash_config.causal is pinned False) + - all 5 layers keep use_k_eq_v=True -> the fused-KV assert still passes + - layer_types stays uniform -> the V2 model runner is NOT required + NOT yet validated end-to-end on GPU; no acceptance-length number exists. + +APPLY + cd && git apply /path/to/gemma4_dspark_swa_vllm.patch + +--- a/vllm/model_executor/models/gemma4_dspark.py ++++ b/vllm/model_executor/models/gemma4_dspark.py +@@ -8,8 +8,9 @@ + import torch.nn as nn + import torch.nn.functional as F + + from vllm import _custom_ops as ops ++from vllm.attention import Attention + from vllm.compilation.decorators import support_torch_compile + from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config + from vllm.model_executor.layers.layernorm import RMSNorm + from vllm.model_executor.layers.linear import ColumnParallelLinear, ReplicatedLinear +@@ -22,9 +23,13 @@ + from vllm.model_executor.model_loader.weight_utils import default_weight_loader + from vllm.transformers_utils.configs.gemma4 import gemma4_layer_config + + from .gemma4_mtp import Gemma4MTPAttention, Gemma4MTPDecoderLayer +-from .qwen3_dflash import DFlashQwen3Model, _dflash_layer_causal ++from .qwen3_dflash import ( ++ DFlashQwen3Model, ++ _dflash_layer_causal, ++ _resolve_layer_attention, ++) + from .qwen3_dspark import DSparkMarkovHead, Qwen3DSparkForCausalLM + from .utils import extract_layer_index, maybe_prefix + + +@@ -61,8 +66,34 @@ + ) + self.is_kv_shared_layer = False + self.causal = _dflash_layer_causal(config, layer_idx) + self.kv_size = self.num_kv_heads * self.head_dim ++ ++ # Gemma4MTPAttention derives the sliding window from the BASE model's ++ # layer pattern (layer_type == "sliding_attention"). That is wrong for a ++ # DFlash/DSpark draft: the draft declares its window via ++ # dflash_config.use_swa / swa_window_size, and deliberately keeps ++ # layer_types all "full_attention" so that attention_k_eq_v stays uniform ++ # for the fused context-KV precompute. Without this, a draft TRAINED with ++ # a sliding window is SERVED with full attention -- no error is raised, ++ # acceptance length just quietly drops. Resolve the window the same way ++ # the Qwen3 DFlash layer does, and rebuild self.attn when it applies. ++ sliding_window, _ = _resolve_layer_attention(config, layer_idx) ++ if sliding_window is not None: ++ get_current_vllm_config().compilation_config.static_forward_context.pop( ++ f"{prefix}.attn", None ++ ) ++ self.attn = Attention( ++ self.num_heads, ++ self.head_dim, ++ self.scaling, ++ num_kv_heads=self.num_kv_heads, ++ cache_config=cache_config, ++ quant_config=quant_config, ++ logits_soft_cap=getattr(config, "attn_logit_softcapping", None), ++ per_layer_sliding_window=sliding_window, ++ prefix=f"{prefix}.attn", ++ ) + attn_bias = getattr(config, "attention_bias", False) + self.k_proj = ColumnParallelLinear( + config.hidden_size, + self.total_num_kv_heads * self.head_dim, diff --git a/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml new file mode 100644 index 00000000000..d87badefad1 --- /dev/null +++ b/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml @@ -0,0 +1,202 @@ +# DSpark recipe for Gemma-4-E4B-it -- SLIDING-WINDOW (SWA) VARIANT. +# +# Derived from dspark_gemma4_e4b.yaml (the validated full-attention baseline); +# the ONLY delta is `dflash_swa_window_size` below. Keep the two in sync. +# +# !! NOT YET SERVABLE -- requires a vLLM fix. Read this before running. !! +# +# Training side is ready: _build_draft_attention_mask() already windows the +# context (`kv > q_real_pos - window`) while leaving block-internal attention +# bidirectional, and hf_dspark.py passes dflash_swa_window_size through. The +# exporter emits dflash_config.use_swa/swa_window_size + a top-level +# sliding_window, and pins causal: False. +# +# The serving gap is Gemma4-specific. vLLM's Qwen3 DFlash layer resolves its +# window via _resolve_layer_attention(), which honours dflash_config.use_swa. +# Gemma4DSparkAttention does NOT: it inherits Gemma4MTPAttention.__init__, +# which derives the window purely from the base-model layer pattern -- +# +# self.is_sliding = layer_type == "sliding_attention" +# sliding_window = config.sliding_window if self.is_sliding else None +# +# Our draft is deliberately all "full_attention" (that is what keeps +# attention_k_eq_v uniform and the fused-KV precompute assert satisfied), so +# is_sliding is False on every layer and per_layer_sliding_window stays None. +# The draft would therefore TRAIN with a 512-token window and SERVE with full +# attention -- a silent train/inference mismatch that only shows up as a +# depressed acceptance length, with no error raised. +# +# Fix (vLLM side, small): have Gemma4DSparkAttention resolve its window through +# _resolve_layer_attention(config, layer_idx) like the Qwen3 DFlash layer does, +# and pass the result as per_layer_sliding_window. Verified by config +# simulation that this keeps all 5 layers use_k_eq_v=True (fused-KV assert +# passes), causal=False, and does not require the V2 model runner. +# +# Until that lands, train with dspark_gemma4_e4b.yaml instead. +# +# +# 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. Note this is the +# LOCAL theta even though the draft below is all full_attention: the +# full_attention entry's "proportional" rope + partial_rotary_factor 0.25 is +# not implemented in the draft classes, and the reference drafter +# (deepseek-ai/dspark_gemma4_12b_block7) also carries the default/1e4 form. +# The draft is trained and served with the same value, so it is self-consistent. +# +# * 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 `` token at id 4 — used directly. +# +# * DRAFT ATTENTION IS NON-CAUSAL (bidirectional), like every other DSpark head. +# vLLM resolves this per layer in qwen3_dflash._dflash_layer_causal: an +# explicit `dflash_config.causal` overrides everything, otherwise a layer is +# causal only when layer_types[i] == "sliding_attention". This draft sets +# neither, and every layer is full_attention, so all 5 layers come out +# causal=False. That is what a block-parallel (semi-autoregressive) drafter +# wants: the block's tokens are predicted at once and must see each other. +# ModelOpt matches this -- DFlashAttention is non-causal, and the exporter +# does not emit a `causal` field. See "SERVING" below for the consequence. +# +# * SERVING requires speculative_config.attention_backend = "FLASHINFER", for +# TWO independent reasons; either alone is fatal: +# 1. head dim -- full_attention layers use global_head_dim=512 and the +# auto-selected FLASH_ATTN (FA2) caps head dimension at 256. +# 2. causality -- load_dspark_model sets attention_config.use_non_causal +# from dflash_has_any_non_causal(), which is True here, and the backend +# must be able to serve a non-causal mask. +# VLLM_ATTENTION_BACKEND does NOT reach the draft; it is read from +# speculative_config.attention_backend. Confirmed at runtime by the log line +# "Using FlashInfer for draft model non-causal attention". + +metadata: + recipe_type: speculative_dflash + description: DSpark for Gemma-4-E4B-it, 5-layer draft with non-causal sliding-window + attention (window 512). Trains today; needs a vLLM fix before it can be served. + +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 + # THE ONLY FUNCTIONAL DELTA vs dspark_gemma4_e4b.yaml. + # Non-causal sliding window over the CONTEXT only; block-internal attention + # stays bidirectional and un-windowed (config.py enforces window >= + # dflash_block_size, so a full block always fits). 512 matches the base + # model's own sliding_window. + dflash_swa_window_size: 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 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: + # Aligned with the official reference drafter deepseek-ai/dspark_gemma4_12b_block7 + # (the checkpoint vLLM's gemma4_dspark.py was written for). Its backbone is + # 5 full-attention layers with attention_k_eq_v. + # + # Why not SWA here: Gemma4DSparkModel._build_fused_kv_buffers asserts + # `all(a.use_k_eq_v for a in layers_attn)`, and Gemma4DSparkAttention only + # sets use_k_eq_v when layer_type == "full_attention". So the assert is + # really about k_eq_v uniformity, but because the two are coupled, ANY + # sliding layer trips it. Mixed sliding/full would additionally need the V2 + # model runner (see _resolve_layer_attention in qwen3_dflash.py). + num_hidden_layers: 5 + num_attention_heads: 16 + num_key_value_heads: 8 + # Full-attention layers use global_head_dim (512) and, under k_eq_v, + # num_global_key_value_heads (1) -- see gemma4_layer_config in vLLM. + head_dim: 256 + global_head_dim: 512 + num_global_key_value_heads: 1 + intermediate_size: 10240 + projector_type: dspark + markov_rank: 256 + markov_head_type: vanilla + use_confidence_head: true + # --- Gemma4 block shape (sandwich norms + layer_scalar) + k_eq_v --- + # model_type gemma4* selects DFlashGemma4DecoderLayer, which adds + # pre/post_feedforward_layernorm and layer_scalar. vLLM's + # Gemma4MTPDecoderLayer looks those up BY NAME and its DSpark loader leaves + # anything it cannot find randomly initialized -- silently. A Qwen3-shaped + # draft therefore loads without error and produces garbage. + model_type: gemma4_text + attention_k_eq_v: true + layer_types: + - full_attention + - full_attention + - full_attention + - full_attention + - full_attention From b27b8858be4b3fdc134ad51510a32e929af72207 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:02:07 +0000 Subject: [PATCH 7/8] example(launcher): Gemma-4-E4B DSpark SWA streaming smoke Single-node co-located streaming smoke for the sliding-window DSpark variant. Derived from hf_streaming_dspark_smoke.yaml; the only deltas are the recipe (dspark_gemma4_e4b_swa.yaml) and the output dir. Training only. The header states plainly that the resulting checkpoint is not servable until the vLLM patch in examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch lands, since Gemma4DSparkAttention would otherwise serve an SWA-trained draft with full attention and no error. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../hf_streaming_dspark_smoke_swa.yaml | 111 ++++++++++++++++++ 1 file changed, 111 insertions(+) create mode 100644 tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml new file mode 100644 index 00000000000..f4cd8f0198c --- /dev/null +++ b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_smoke_swa.yaml @@ -0,0 +1,111 @@ +# DSpark streaming SMOKE run (co-located single node) for Gemma-4-E4B-it on +# AWS-PDX -- SLIDING-WINDOW (SWA) VARIANT. +# +# Derived from hf_streaming_dspark_smoke.yaml; the only deltas are the recipe +# (dspark_gemma4_e4b_swa.yaml, which sets dflash_swa_window_size: 512) and the +# output dir. Keep the two in sync. +# +# TRAINING ONLY. The resulting checkpoint is NOT servable as-is: vLLM's +# Gemma4DSparkAttention derives its sliding window from the BASE layer pattern +# instead of dflash_config.use_swa, so an SWA-trained draft would be served with +# FULL attention -- silently, showing up only as depressed acceptance. See +# examples/speculative_decoding/export/gemma4_dspark_swa_vllm.patch. The point of +# this run is to prove SWA training converges and to produce a checkpoint. +# +# +# Purpose: prove the streaming chain end-to-end (vLLM serve -> aux hidden-state +# capture -> NIXL transfer -> fake-base trainer) for a NEW base model. Only ~20 +# steps on a 1024-row corpus; the numbers are meaningless, the point is that the +# pipeline runs and the loss is finite and decreasing. +# +# Topology: serve TP=1 on GPU 0 (E4B is 16 GB, fits one B300), DSpark trainer on +# GPUs 4-7. Intra-node NIXL, no cross-node EFA. +# +# Gemma-4-E4B specifics verified on PDX 2026-08-12: +# * EAGLE_CAPTURE_IDS: vLLM's EagleModelMixin captures POST-layer with the +# residual added and indexes as layer_idx+1, so valid ids are 1..42 and 42 is +# the TRUE final layer (verified cos=1.0000 against HF hidden_states, with +# off-by-one dropping to 0.55-0.98). Gemma 4 repeats 5x sliding + 1x full +# attention, so full-attention layers sit at capture ids +# [6,12,18,24,30,36,42]; we sample those to land on residual-stream +# boundaries instead of spacing uniformly (deep layers are near-redundant: +# adjacent cosine ~0.98-0.99 around id 36 vs 0.55-0.70 around id 9). +# * NO --trust-remote-code needed (the repo ships no .py) and no vLLM patch: +# gemma4 is natively supported in the 2026-08-11 nightly container. +# * Corpus has its user-only `messages` column DROPPED — hf_streaming_dataset +# prefers `messages` over `conversations` and a user-only one makes streaming +# SILENTLY HANG. + +job_name: Gemma-4-E4B_DSpark_streaming_smoke_swa +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/gemma-4-E4B-it + + task_0: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b_swa.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.mode=streaming + # Gemma 4's stock chat template has NO generation markers, so + # answer_only_loss silently yields an all-zero loss_mask and EVERY row is + # rejected ("no fetchable sample found in the entire corpus"). This copy + # adds them; see the template file header for details. + - data.chat_template=examples/google/gemma-4-E4B-it/chat_template_train.jinja + - data.data_path=/smokedata + - training.output_dir=/scratchspace/dspark_smoke_swa + - training.training_seq_len=2048 + - training.disable_tqdm=true + - training.ar_validate_steps=500000 + - training.num_train_epochs=1 + - training.max_steps=20 + - training.logging_steps=1 + - training.save_steps=20 + - training.per_device_train_batch_size=1 + - training.answer_only_loss=true + - training.report_to=none + environment: + - HF_MODEL_CKPT: <> + # MUST be exactly num_draft_layers + 1 entries (5 draft layers -> 6 ids): + # the projector is sized from the DRAFT's num_hidden_layers, so 7 ids gave + # "mat1 and mat2 shapes cannot be multiplied (2048x15360 and 12800x2560)". + # Confirmed against the working runs: K2.6 6 layers -> 7 ids, gpt-oss 5 -> 6. + # Chosen from the full-attention layers of the 5:1 sliding/full cycle + # ([6,12,18,24,30,36,42]), dropping 30 to keep both ends and the true final 42. + - EAGLE_CAPTURE_IDS: "[6,12,18,24,36,42]" + - SERVE_GPU: "0" + - SERVE_TP: "1" + - SERVE_GPU_MEM_UTIL: "0.85" + - STREAMING_NUM_WORKERS: "1" + - SERVE_MAX_MODEL_LEN: "2176" + - SERVE_MAX_NUM_SEQS: "4" + - SERVE_READY_TIMEOUT: "2400" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + # NIXL hidden-state transport. Without these, NIXL falls back to the UCX + # backend, which reports "8 NVIDIA GPU(s) were detected, but UCX CUDA + # support was not found" and then dies with NIXL_ERR_REMOTE_DISCONNECT the + # moment the trainer tries to pull hidden states. LIBFABRIC+efa is what the + # working Kimi-K2.6 streaming runs use. + - NIXL_BACKENDS: LIBFABRIC + - FI_PROVIDER: efa + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 + # The 2026-08-11 nem35 image has NO libfabric and no NIXL plugins, so + # NIXL_BACKENDS=LIBFABRIC dies with NIXL_ERR_NOT_FOUND and the UCX fallback + # dies with NIXL_ERR_REMOTE_DISCONNECT ("UCX CUDA support was not found"). + # The auxfix image carries the AWS libfabric stack and is what the working + # Kimi-K2.6 streaming runs use. Its vLLM is older (transformers 5.12.1) -- + # gemma4 support must be re-verified in THIS image. + container: /home/haoguo/lustre/containers/vllm-nightly-efa-x86_64-auxfix.sqsh + container_mounts: + - /home/haoguo/lustre/hf-local:/hf-local + - /home/haoguo/lustre/g4_smoke_corpus:/smokedata From e82d1cad2e89363e9057cd70e504650f4b1eaa71 Mon Sep 17 00:00:00 2001 From: h-guo18 <67671475+h-guo18@users.noreply.github.com> Date: Wed, 19 Aug 2026 12:15:39 +0000 Subject: [PATCH 8/8] example(launcher): Gemma-4-E4B DSpark full 2-epoch streaming run The 4-node launcher config for the first real (non-smoke) DSpark training run on Gemma-4-E4B-it, over a 1.13M-row synthesized corpus. Topology is 2 serve + 2 trainer nodes: the launcher splits WHOLE nodes for nodes >= 2, so 2:2 is what yields a clean global batch of 64 (16 DP ranks x per_device_train_batch_size 4 x grad_accum 1). The header records two traps that cost real debugging time: * the corpus MUST be the cleaned copy -- the raw synthesis output carries a user-only `messages` column, and hf_streaming_dataset prefers `messages` over `conversations`, so answer_only_loss yields an all-zero loss mask, every row is rejected, and streaming SILENTLY HANGS with no error; * the auxfix container is required for NIXL's libfabric transport -- the nem35 image has no libfabric, so NIXL falls back to UCX and dies with NIXL_ERR_REMOTE_DISCONNECT as soon as the trainer pulls hidden states. Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> --- .../hf_streaming_dspark_full.yaml | 96 +++++++++++++++++++ 1 file changed, 96 insertions(+) create mode 100644 tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml diff --git a/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml new file mode 100644 index 00000000000..12b8ea36ff0 --- /dev/null +++ b/tools/launcher/examples/google/gemma-4-E4B-it/hf_streaming_dspark_full.yaml @@ -0,0 +1,96 @@ +# DSpark FULL training run for Gemma-4-E4B-it on AWS-PDX. +# +# First real training run (the earlier 20-step smoke was a pipeline test only). +# Corpus: 1,130,852 synthesized rows from ~/lustre/g4_corpus_train. +# +# TOPOLOGY: 4 nodes on the interactive QOS (its hard cap), split 2 serve + 2 +# trainer. The launcher does NOT support "own serve on every node" -- for +# nodes >= 2 it splits WHOLE nodes (see train_eagle_streaming.sh header), so +# 2:2 is what yields a clean global batch of 64: +# +# 2 trainer nodes x 8 GPU = 16 DP ranks +# per_device_train_batch_size 4 x grad_accum 1 +# -> global batch = 16 * 4 * 1 = 64 +# +# 2 epochs: 1,130,852 * 2 / 64 ~= 35,340 steps +# +# CORPUS: must be the CLEANED copy (g4_corpus_train). The raw synthesis output +# carries a user-only `messages` column, and hf_streaming_dataset prefers +# `messages` over `conversations` -- that yields an all-zero loss mask under +# answer_only_loss, every row is rejected, and streaming SILENTLY HANGS with no +# error. Verified 200/200 raw rows were user-only; the cleaned copy has 0. +# +# Gemma-4-E4B specifics (verified on PDX 2026-08-12, unchanged): +# * EAGLE_CAPTURE_IDS: vLLM captures POST-layer with the residual added and +# indexes as layer_idx+1, so valid ids are 1..42 and 42 is the TRUE final +# layer. Gemma 4 repeats 5x sliding + 1x full attention, so full-attention +# layers sit at ids [6,12,18,24,30,36,42]; we sample 6 of those (must be +# num_draft_layers + 1 = 6) to land on residual-stream boundaries. +# * No --trust-remote-code needed and no vLLM patch; gemma4 is native in the +# 2026-08-11 nightly. + +job_name: Gemma-4-E4B_DSpark_full_2ep +pipeline: + allow_to_fail: false + skip: false + note: + + global_vars: + hf_model: /hf-local/gemma-4-E4B-it + + task_0: + script: common/eagle3/train_eagle_streaming.sh + args: + - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark_gemma4_e4b.yaml + - model.model_name_or_path=<> + - model.use_fake_base_for_offline=true + - data.mode=streaming + # Gemma 4's stock chat template has NO generation markers, so + # answer_only_loss silently yields an all-zero loss_mask and every row is + # rejected. This copy adds them; see the template file header. + - data.chat_template=examples/google/gemma-4-E4B-it/chat_template_train.jinja + - data.data_path=/traindata + - training.output_dir=/scratchspace/dspark_full_2ep + - training.training_seq_len=4096 + - training.disable_tqdm=true + - training.ar_validate_steps=0 + - training.num_train_epochs=2 + - training.per_device_train_batch_size=4 + - training.gradient_accumulation_steps=1 + - training.learning_rate=1.0e-4 + - training.warmup_steps=500 + - training.logging_steps=20 + - training.save_steps=1000 + # All checkpoints are kept (user's call): HF save_total_limit defaults to + # None, which retains every one, so no override is needed. + - training.answer_only_loss=true + - training.report_to=none + environment: + - HF_MODEL_CKPT: <> + # MUST be exactly num_draft_layers + 1 (5 draft layers -> 6 ids). + - EAGLE_CAPTURE_IDS: "[6,12,18,24,36,42]" + # 2 serve nodes + 2 trainer nodes (of the 4 allocated). + - SERVE_NODES: "2" + - SERVE_GPU_MEM_UTIL: "0.9" + - SERVE_MAX_MODEL_LEN: "4352" + - SERVE_MAX_NUM_SEQS: "64" + - SERVE_READY_TIMEOUT: "2400" + - STREAMING_NUM_WORKERS: "4" + - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200" + - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200" + # NIXL hidden-state transport. Without these NIXL falls back to UCX, which + # reports "UCX CUDA support was not found" and dies with + # NIXL_ERR_REMOTE_DISCONNECT the moment the trainer pulls hidden states. + - NIXL_BACKENDS: LIBFABRIC + - FI_PROVIDER: efa + slurm_config: + _factory_: "slurm_factory" + nodes: 4 + ntasks_per_node: 1 + gpus_per_node: 8 + # The auxfix image carries the AWS libfabric stack that NIXL needs; the + # nem35 image has no libfabric and no NIXL plugins. + container: /home/haoguo/lustre/containers/vllm-nightly-efa-x86_64-auxfix.sqsh + container_mounts: + - /home/haoguo/lustre/hf-local:/hf-local + - /home/haoguo/lustre/g4_corpus_train:/traindata