nemo_automodel.components.models.gemma4_drafter.model

Thin NeMo wrapper around HuggingFace Gemma4AssistantForCausalLM.

The HF implementation in transformers.models.gemma4_assistant is used as-is, except that the ordered-embedding output head scores the full vocabulary (see Gemma4DrafterFullVocabEmbedder). The wrapper adds HFCheckpointingMixin so the drafter participates in NeMo’s distributed checkpointing pipeline and gives us a stable native class name for the model registry.

Module Contents

Classes

Name	Description
`Gemma4DrafterForCausalLM`	Placeholder raised when `transformers.models.gemma4_assistant` is unavailable.
`Gemma4DrafterFullVocabEmbedder`	Ordered-embedding head that scores the full vocabulary.

Data

ModelClass

_GEMMA4_ASSISTANT_HF_AVAILABLE

_GEMMA4_ASSISTANT_MODELING

_GEMMA4_ASSISTANT_MSG

API

class nemo_automodel.components.models.gemma4_drafter.model.Gemma4DrafterForCausalLM(
    args = (),
    kwargs = {}
)

Placeholder raised when transformers.models.gemma4_assistant is unavailable.

nemo_automodel.components.models.gemma4_drafter.model.Gemma4DrafterForCausalLM.from_pretrained(
    args = (),
    kwargs = {}
)

classmethod

class nemo_automodel.components.models.gemma4_drafter.model.Gemma4DrafterFullVocabEmbedder()

Bases: HFGemma4AssistantMaskedEmbedder

Ordered-embedding head that scores the full vocabulary.

The HF ordered head scores only the tokens of the top-k centroid clusters and gives every other token a constant logit without gradient, so training through it diverges (issue #4022). This head keeps the centroids and token_ordering state of the HF head, so checkpoints and exports are unchanged, but always returns full-vocabulary logits, in both training and eval mode. Only an exported checkpoint reloaded by HF or vLLM uses the ordered head.

nemo_automodel.components.models.gemma4_drafter.model.Gemma4DrafterFullVocabEmbedder.forward(
    hidden_states: torch.Tensor,
    lm_head_weight: torch.Tensor
) -> torch.Tensor

Score every vocabulary token.

Parameters:

hidden_states

torch.Tensor

Drafter hidden states of shape [batch, seq_len, hidden_size].

lm_head_weight

torch.Tensor

Tied output embedding of shape [vocab_size, hidden_size].

Returns: torch.Tensor

Logits of shape [batch, seq_len, vocab_size], hidden_states @ lm_head_weight.T.

nemo_automodel.components.models.gemma4_drafter.model.ModelClass = Gemma4DrafterForCausalLM

nemo_automodel.components.models.gemma4_drafter.model._GEMMA4_ASSISTANT_HF_AVAILABLE = _HAS_CONFIG and _HAS_CAUSAL_LM and _HAS_MASKED_EMBEDDER

nemo_automodel.components.models.gemma4_drafter.model._GEMMA4_ASSISTANT_MODELING = 'transformers.models.gemma4_assistant.modeling_gemma4_assistant'

nemo_automodel.components.models.gemma4_drafter.model._GEMMA4_ASSISTANT_MSG = 'transformers.models.gemma4_assistant is not available.'