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