> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# 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`](#nemo_automodel-components-models-gemma4_drafter-model-Gemma4DrafterForCausalLM)             | Placeholder raised when `transformers.models.gemma4_assistant` is unavailable. |
| [`Gemma4DrafterFullVocabEmbedder`](#nemo_automodel-components-models-gemma4_drafter-model-Gemma4DrafterFullVocabEmbedder) | Ordered-embedding head that scores the full vocabulary.                        |

### Data

[`ModelClass`](#nemo_automodel-components-models-gemma4_drafter-model-ModelClass)

[`_GEMMA4_ASSISTANT_HF_AVAILABLE`](#nemo_automodel-components-models-gemma4_drafter-model-_GEMMA4_ASSISTANT_HF_AVAILABLE)

[`_GEMMA4_ASSISTANT_MODELING`](#nemo_automodel-components-models-gemma4_drafter-model-_GEMMA4_ASSISTANT_MODELING)

[`_GEMMA4_ASSISTANT_MSG`](#nemo_automodel-components-models-gemma4_drafter-model-_GEMMA4_ASSISTANT_MSG)

### API

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

Placeholder raised when `transformers.models.gemma4_assistant` is unavailable.

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

classmethod

```python
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.

```python
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`.

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

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

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

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