> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/gemma4_drafter/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

Gemma4 drafter (assistant) NeMo Automodel support.

Adds a thin NeMo wrapper around the HuggingFace `Gemma4AssistantForCausalLM`
plus a `Gemma4WithDrafter` composite for joint base + drafter fine-tuning.

## Submodules

* **[`nemo_automodel.components.models.gemma4_drafter.composite`](/nemo-automodel/nemo_automodel/components/models/gemma4_drafter/composite)**
* **[`nemo_automodel.components.models.gemma4_drafter.model`](/nemo-automodel/nemo_automodel/components/models/gemma4_drafter/model)**

## Package Contents

### Classes

| Name                                                                                                          | Description                                                                    |
| ------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`Gemma4DrafterForCausalLM`](#nemo_automodel-components-models-gemma4_drafter-model-Gemma4DrafterForCausalLM) | Placeholder raised when `transformers.models.gemma4_assistant` is unavailable. |
| [`Gemma4JointOutput`](#nemo_automodel-components-models-gemma4_drafter-composite-Gemma4JointOutput)           | Output of `Gemma4WithDrafter`.                                                 |
| [`Gemma4WithDrafter`](#nemo_automodel-components-models-gemma4_drafter-composite-Gemma4WithDrafter)           | Composite model that wraps a Gemma 4 base + its released drafter.              |

### Data

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

### 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.composite.Gemma4JointOutput(
    logits: torch.Tensor,
    drafter_logits: list[torch.Tensor] = list(),
    drafter_loss_weight: float = 1.0,
    hidden_states: tuple | None = None,
    loss: torch.Tensor | None = None
)
```

Dataclass

Output of `Gemma4WithDrafter`.

**`drafter_logits`** `list[Tensor] = field(default_factory=list)`

Per-step list of drafter logits, each `[B, S, V]`.
For the default single-step recurrent cell this list has length 1.

---

**`drafter_loss_weight`** `float = 1.0`

`lambda` multiplier the recipe applies to the
drafter loss when summing it with the base loss.

---

**`hidden_states`** `tuple | None = None`

Optional list of base hidden states (mirrors HF).

---

**`logits`** `Tensor`

Base model logits `[B, S, V]`.

---

**`loss`** `Tensor | None = None`

Placeholder, populated by the recipe if needed.

---

```python
class nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter(
    base: torch.nn.Module,
    drafter: torch.nn.Module,
    drafter_loss_weight: float = 1.0,
    drafter_num_steps: int = 1,
    freeze_base_for_drafter: bool = False,
    share_embedding_with_base: bool = False,
    base_activation_checkpointing: bool = False
)
```

**Bases:** `Module`, [HFCheckpointingMixin](/nemo-automodel/nemo_automodel/components/models/common/hf_checkpointing_mixin#nemo_automodel-components-models-common-hf_checkpointing_mixin-HFCheckpointingMixin)

Composite model that wraps a Gemma 4 base + its released drafter.

Both sub-modules are loaded via NeMo's `NeMoAutoModel*` paths so they
receive the standard distributed infrastructure (FSDP2 sharding, freeze
config, checkpoint loading, kernel patches, ...) independently. The
composite is a thin `nn.Module` that owns both and exposes a joint
forward and a `save_pretrained` that writes the pair as two HF-format
sub-directories (`base/` and `drafter/`).

**Parameters:**

**`base`** `nn.Module`

Loaded base model (typically a `Gemma4ForConditionalGeneration`
instance returned by `NeMoAutoModelForImageTextToText.from_pretrained`).

---

**`drafter`** `nn.Module`

Loaded drafter (a `Gemma4DrafterForCausalLM` instance
returned by `NeMoAutoModelForCausalLM.from_pretrained`).

---

**`drafter_loss_weight`** `float` — default: 1.0

Multiplier `lambda` applied to the drafter loss
in the recipe.

---

**`drafter_num_steps`** `int` — default: 1

Number of recurrent drafter steps K to run per
training batch. With K = 1 the composite is the EAGLE-1-style
single-step setup; with K > 1 the drafter runs autoregressively
for K rounds, feeding its previous-round `last_hidden_state`
(already post-projected to H\_b) and a teacher-forced shifted
token id back into itself, matching the Gemma 4 drafter blog's
recipe. `shared_kv_states` is captured from a single base
forward and reused at every round.

---

**`base_activation_checkpointing`** `= bool(base_activation_checkpointing)`

---

**`drafter_loss_weight`** `= float(drafter_loss_weight)`

---

**`drafter_num_steps`** `= int(drafter_num_steps)`

---

**`freeze_base_for_drafter`** `= bool(freeze_base_for_drafter)`

---

**`share_embedding_with_base`** `= bool(share_embedding_with_base)`

---

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter._get_base_text_config(
    base: torch.nn.Module
)
```

staticmethod

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.forward(
    input_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> nemo_automodel.components.models.gemma4_drafter.composite.Gemma4JointOutput
```

Joint forward: base first, then drafter consuming the base's outputs.

Any extra kwargs (`pixel_values`, `mm_token_type_ids`,
`pixel_values_videos`, `input_features`, ...) are passed straight
through to the base. Multimodal kwargs are *not* forwarded to the
drafter (the drafter is text-only).

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.from_pretrained(
    base_path: str | None = None,
    drafter_path: str | None = None,
    pretrained_model_name_or_path: str | None = None,
    drafter_loss_weight: float = 1.0,
    drafter_num_steps: int = 1,
    freeze_base_for_drafter: bool = False,
    share_embedding_with_base: bool = False,
    base_activation_checkpointing: bool = False,
    torch_dtype: typing.Any = None,
    attn_implementation: str | None = None,
    use_liger_kernel: bool | None = None,
    use_sdpa_patching: bool | None = None,
    text_config: dict | None = None,
    peft_config: typing.Any = None,
    device_mesh: typing.Any = None,
    moe_mesh: typing.Any = None,
    distributed_config: typing.Any = None,
    pipeline_config: typing.Any = None,
    distributed_setup: typing.Any = None,
    freeze_config: typing.Any = None,
    cache_dir: str | None = None,
    kwargs = {}
) -> 'Gemma4WithDrafter'
```

classmethod

Build the composite by loading base and drafter via the NeMoAuto paths.

**Parameters:**

**`base_path`** `str | None` — default: None

HF repo id or local path of the Gemma 4 base model.

---

**`drafter_path`** `str | None` — default: None

HF repo id or local path of the released drafter.

---

**`pretrained_model_name_or_path`** `str | None` — default: None

Alias for `base_path`. Kept so that
YAML configs can set `pretrained_model_name_or_path` and have
the recipe's processor / checkpoint-config helpers (which read
this key from the model config) keep working.

---

**`drafter_loss_weight`** `float` — default: 1.0

`lambda` multiplier on the drafter loss.

---

**`drafter_num_steps`** `int` — default: 1

Number of recurrent drafter steps K per batch.
`K = 1` is EAGLE-1-style single-step; `K > 1` matches the
Gemma 4 drafter blog's multi-token-prediction (MTP) training
recipe -- the drafter consumes its previous round's
post-projected hidden state plus a teacher-forced shifted
token id at every subsequent round.

---

**`freeze_base_for_drafter`** `bool` — default: False

If True, freeze all base parameters so
only the drafter is trained (drafter-only sub-case). Default
False (joint training).

---

**`share_embedding_with_base`** `bool` — default: False

If True, copy the base's input
embedding into the drafter's `embed_tokens` once at init.
The drafter's `lm_head` is tied to its own `embed_tokens`
so the row weights start aligned with the base too. The two
embeddings then evolve as independent parameters during
training.

---

**`base_activation_checkpointing`** `bool` — default: False

If True, enable HF gradient
checkpointing on the base to reduce activation memory.
Important for the 4B + drafter + long-context setting.

---

**`torch_dtype`** `Any` — default: None

dtype to use for both sub-models. Must be
`torch.bfloat16` -- the drafter is bf16-only.

---

**`attn_implementation`** `str | None` — default: None

Forwarded to both sub-loads.

---

**`use_liger_kernel`** `bool | None` — default: None

Forwarded to both sub-loads.

---

**`use_sdpa_patching`** `bool | None` — default: None

Forwarded to both sub-loads.

---

**`text_config`** `dict | None` — default: None

Optional overrides forwarded to the base load.

---

**`peft_config`** `Any` — default: None

PEFT config (currently expected to be `None` --
joint drafter PEFT is out of scope for the initial recipe).

---

**`device_mesh`** `Any` — default: None

Distributed device mesh shared by base and drafter.

---

**`moe_mesh`** `Any` — default: None

MoE mesh shared by base and drafter (drafter is dense).

---

**`distributed_config`** `Any` — default: None

FSDP2 / Megatron-FSDP / DDP config object.

---

**`pipeline_config`** `Any` — default: None

Must be `None` -- pipeline parallelism is not
supported when the drafter is attached.

---

**`distributed_setup`** `Any` — default: None

Resolved `DistributedSetup` (topology + policy)
shared by base and drafter. This is the path used by the VLM
finetune recipe; its `pp_size` and `cp_size` must be `1`.

---

**`freeze_config`** `Any` — default: None

Forwarded to the base only (the drafter is trained
end-to-end). Customize the drafter's freezing with explicit
`requires_grad_` calls on the returned composite if needed.

---

**`cache_dir`** `str | None` — default: None

HuggingFace cache directory.

---

**`**kwargs`** — default: \{}

Additional kwargs forwarded to both sub-loads.

---

**Returns:** `'Gemma4WithDrafter'`

An instantiated `Gemma4WithDrafter`.

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.get_input_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.get_output_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.get_rope_index(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.load_pretrained(
    load_directory: str,
    checkpointer: 'Checkpointer' | None = None,
    kwargs = {}
) -> None
```

Load weights from the two-subdir layout written by `save_pretrained`.

Mirrors the save side: reads `<load_directory>/base/model` and
`<load_directory>/drafter/model` (the standard `Checkpointer.save_model`
output layout) and routes them to `self.base` and `self.drafter`
respectively. Used by the recipe's resume path when a checkpoint
directory was produced by this composite.

**Parameters:**

**`load_directory`** `str`

A checkpoint directory containing `base/` and
`drafter/` sub-directories (e.g. `<ckpt_dir>/epoch_X_step_Y`).

---

**`checkpointer`** `'Checkpointer' | None` — default: None

The recipe's `Checkpointer` instance.

---

**`**kwargs`** — default: \{}

Reserved; ignored.

---

```python
nemo_automodel.components.models.gemma4_drafter.composite.Gemma4WithDrafter.save_pretrained(
    save_directory: str,
    checkpointer: 'Checkpointer' | None = None,
    tokenizer: typing.Any = None,
    kwargs = {}
) -> None
```

Save base and drafter as two HF-format sub-directories.

Produces `<save_directory>/base/` and `<save_directory>/drafter/`
with HF-compatible artifacts. Each side can later be loaded back by HF
`from_pretrained` independently (vLLM compatibility).

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

## Docs

- [nemo_automodel.components.models.gemma4_drafter.composite](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/gemma4_drafter/composite.md)
- [nemo_automodel.components.models.gemma4_drafter.model](https://docs.nvidia.com/nemo/automodel/v0.5/nemo-automodel/nemo_automodel/components/models/gemma4_drafter/model.md)