nemo_automodel.components.speculative.eagle.draft_gemma

EAGLE-3 draft model for Gemma4 targets (Gemma4ForConditionalGeneration).

Gemma4 is a multimodal decoder whose text backbone differs from a Llama-style dense LLM in a handful of ways: a scaled word embedding (* sqrt(hidden)), zero-centred (1 + w) RMSNorm, per-head Q/K norm, GeGLU (gelu_pytorch_tanh) MLPs, alternating sliding-window / full attention with two separate RoPE schedules, and (on the E2B/E4B checkpoints) Gemma-3n-style per-layer inputs and AltUp. Almost none of that reaches the EAGLE-3 draft: the draft is a single from-scratch decoder layer that consumes only the post-block auxiliary hidden states emitted by the frozen target (via register_forward_hook) and re-projects its own Q/K/V, so it never sees the target’s sliding mask, per-layer inputs, AltUp, or experts — structurally it is the same Llama-style dense draft used for every other registry entry. The draft’s own RMSNorms are trained from scratch, so the target’s zero-centred norm form does not need to be reproduced, and the constant embedding scale is normalised away by the draft’s input_layernorm before the fused [embed, hidden] attention input.

Three config quirks do have to be reconciled before the shared Llama draft can build, and that is all this module does (see :func:_normalize_gemma4_draft_config):

1. Activation key. Gemma text configs name the MLP activation hidden_activation (gelu_pytorch_tanh); the shared MLP reads config.hidden_act. The GeGLU structure is identical to the draft’s SwiGLU wiring once the activation is gelu_pytorch_tanh, so we copy the value across.
2. RoPE. Gemma4 stores a nested, per-attention-type rope_parameters (full_attention: rope_theta=1e6, rope_type="proportional", partial_rotary_factor=0.25; sliding_attention: rope_theta=1e4, rope_type="default") that the shared LlamaRotaryEmbedding cannot read — its _get_rope_config expects a flat rope_theta and would silently fall back to base 10000. The single draft layer runs full causal attention over the whole sequence, so it mirrors the target’s global (full-attention) schedule. We flatten that schedule to a standard full-rotary Llama RoPE (rope_theta = the global theta, head_dim from the config, no partial rotary). Training and inference are then self-consistent: the saved checkpoint keeps the canonical architectures: ["LlamaEagle3DraftModel"] string and a flat rope_theta, so SGLang / vLLM reproduce the exact same rotary with their existing EAGLE-3 Llama head — no Gemma-specific inference support is required. (Reproducing the target’s proportional + partial-rotary global schedule instead would be more faithful to the target’s long-context frequencies but no inference engine can currently serve it, so it is intentionally not done here.)
3. MoE FFN width. On a MoE Gemma4 (enable_moe_block) the config intermediate_size is the per-expert FFN width (e.g. 2112 on 26B-A4B, below hidden_size 2816). The dense draft would otherwise build a contracting MLP and be starved of capacity, so the draft MLP is sized to the target’s active FFN width (top_k_experts * intermediate_size). Dense targets are untouched.

Everything else (GQA, the EAGLE-3 TTT cache attention, the fc projection, the draft lm_head and vocab mapping) is inherited unchanged from :class:LlamaEagle3DraftModel, and the on-disk state-dict layout is identical.

Module Contents

Classes

Functions

Data

_DEFAULT_GLOBAL_ROPE_THETA

logger

API

class nemo_automodel.components.speculative.eagle.draft_gemma.Gemma4Eagle3DraftModel(
    config: transformers.PretrainedConfig
)

nemo_automodel.components.speculative.eagle.draft_gemma._extract_global_rope_theta(
    config: transformers.PretrainedConfig
) -> float

nemo_automodel.components.speculative.eagle.draft_gemma._normalize_gemma4_draft_config(
    config: transformers.PretrainedConfig
) -> None

nemo_automodel.components.speculative.eagle.draft_gemma._DEFAULT_GLOBAL_ROPE_THETA = 1000000.0

nemo_automodel.components.speculative.eagle.draft_gemma.logger = logging.getLogger(__name__)