> 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.bagel

BAGEL mixed-modal LLM (Qwen2 + SigLIP-NaViT + MoT shell).

## Submodules

* **[`nemo_automodel.components.models.bagel.attention_masks`](/nemo-automodel/nemo_automodel/components/models/bagel/attention_masks)**
* **[`nemo_automodel.components.models.bagel.autoencoder`](/nemo-automodel/nemo_automodel/components/models/bagel/autoencoder)**
* **[`nemo_automodel.components.models.bagel.configuration`](/nemo-automodel/nemo_automodel/components/models/bagel/configuration)**
* **[`nemo_automodel.components.models.bagel.connector`](/nemo-automodel/nemo_automodel/components/models/bagel/connector)**
* **[`nemo_automodel.components.models.bagel.embeddings`](/nemo-automodel/nemo_automodel/components/models/bagel/embeddings)**
* **[`nemo_automodel.components.models.bagel.hf_backbone_loader`](/nemo-automodel/nemo_automodel/components/models/bagel/hf_backbone_loader)**
* **[`nemo_automodel.components.models.bagel.model`](/nemo-automodel/nemo_automodel/components/models/bagel/model)**
* **[`nemo_automodel.components.models.bagel.modeling_qwen2_packed`](/nemo-automodel/nemo_automodel/components/models/bagel/modeling_qwen2_packed)**
* **[`nemo_automodel.components.models.bagel.modeling_siglip_navit`](/nemo-automodel/nemo_automodel/components/models/bagel/modeling_siglip_navit)**
* **[`nemo_automodel.components.models.bagel.parallelization`](/nemo-automodel/nemo_automodel/components/models/bagel/parallelization)**
* **[`nemo_automodel.components.models.bagel.state_dict_adapter`](/nemo-automodel/nemo_automodel/components/models/bagel/state_dict_adapter)**

## Package Contents

### Classes

| Name                                                                                                   | Description                                                              |
| ------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------ |
| [`AutoEncoder`](#nemo_automodel-components-models-bagel-autoencoder-AutoEncoder)                       | BAGEL Stage 2 autoencoder wrapper.                                       |
| [`AutoEncoderParams`](#nemo_automodel-components-models-bagel-autoencoder-AutoEncoderParams)           | Architecture parameters for the BAGEL/FLUX autoencoder.                  |
| [`BagelConfig`](#nemo_automodel-components-models-bagel-configuration-BagelConfig)                     | Top-level BAGEL config.                                                  |
| [`BagelForUnifiedMultimodal`](#nemo_automodel-components-models-bagel-model-BagelForUnifiedMultimodal) | BAGEL mixed-modal LLM wrapper for understanding and optional generation. |

### Functions

| Name                                                                                                   | Description                                       |
| ------------------------------------------------------------------------------------------------------ | ------------------------------------------------- |
| [`load_bagel_autoencoder`](#nemo_automodel-components-models-bagel-autoencoder-load_bagel_autoencoder) | Load the BAGEL autoencoder from `ae.safetensors`. |

### API

```python
class nemo_automodel.components.models.bagel.autoencoder.AutoEncoder(
    params: nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams
)
```

**Bases:** `Module`

BAGEL Stage 2 autoencoder wrapper.

**`decoder`**

---

**`encoder`**

---

**`reg`** `= DiagonalGaussian()`

---

**`scale_factor`** `= params.scale_factor`

---

**`shift_factor`** `= params.shift_factor`

---

```python
nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.decode(
    z: torch.Tensor
) -> torch.Tensor
```

Decode scaled latents to image tensors.

```python
nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.encode(
    x: torch.Tensor
) -> torch.Tensor
```

Encode image tensors to scaled latents.

```python
nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.forward(
    x: torch.Tensor
) -> torch.Tensor
```

Encode and decode image tensors.

```python
class nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams(
    resolution: int,
    in_channels: int,
    downsample: int,
    ch: int,
    out_ch: int,
    ch_mult: list[int],
    num_res_blocks: int,
    z_channels: int,
    scale_factor: float,
    shift_factor: float
)
```

Dataclass

Architecture parameters for the BAGEL/FLUX autoencoder.

**`ch`** `int`

---

**`ch_mult`** `list[int]`

---

**`downsample`** `int`

---

**`in_channels`** `int`

---

**`num_res_blocks`** `int`

---

**`out_ch`** `int`

---

**`resolution`** `int`

---

**`scale_factor`** `float`

---

**`shift_factor`** `float`

---

**`z_channels`** `int`

---

```python
class nemo_automodel.components.models.bagel.configuration.BagelConfig(
    vision_config: typing.Union[typing.Dict[str, typing.Any], nemo_automodel.components.models.bagel.modeling_siglip_navit.SiglipVisionConfig, None] = None,
    text_config: typing.Union[typing.Dict[str, typing.Any], transformers.Qwen2Config, None] = None,
    visual_und: bool = True,
    visual_gen: bool = False,
    stage: typing.Union[int, str, None] = None,
    llm_path: str = '',
    vit_path: str = '',
    vae_path: str = '',
    max_latent_size: int = 64,
    latent_patch_size: int = 2,
    vit_patch_size: int = 14,
    vit_max_num_patch_per_side: int = 70,
    connector_act: str = 'gelu_pytorch_tanh',
    interpolate_pos: bool = False,
    vit_select_layer: int = -2,
    vit_rope: bool = False,
    text_cond_dropout_prob: float = 0.1,
    vae_cond_dropout_prob: float = 0.3,
    vit_cond_dropout_prob: float = 0.3,
    timestep_shift: float = 1.0,
    pad_token_id: int = 151643,
    llm_config: typing.Union[typing.Dict[str, typing.Any], transformers.Qwen2Config, None] = None,
    vit_config: typing.Union[typing.Dict[str, typing.Any], nemo_automodel.components.models.bagel.modeling_siglip_navit.SiglipVisionConfig, None] = None,
    vae_config: typing.Union[typing.Dict[str, typing.Any], None] = None,
    kwargs: typing.Any = {}
)
```

**Bases:** `PretrainedConfig`

Top-level BAGEL config.

The text and vision sub-configs are nested `PretrainedConfig`
instances (not bare dicts) so callers can introspect them the same way
they would with any other HF config.

Attribute aliases:

* `text_config` \<-> `llm_config` (both point at the same object)
* `vision_config` \<-> `vit_config` (ditto)

**`architectures`** `= ['BagelForUnifiedMultimodal']`

---

**`llm_config`** `Qwen2Config`

---

**`model_type`** `= 'bagel'`

---

**`text_config`** `= _coerce_text_config(text_config)`

---

**`vae_config`** `= vae_config or {}`

---

**`vision_config`** `= _coerce_vision_config(vision_config)`

---

**`vit_config`** `SiglipVisionConfig | None`

---

```python
nemo_automodel.components.models.bagel.configuration.BagelConfig.to_dict() -> typing.Dict[str, typing.Any]
```

```python
class nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal(
    config: nemo_automodel.components.models.bagel.configuration.BagelConfig,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

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

BAGEL mixed-modal LLM wrapper for understanding and optional generation.

`visual_gen=False` gives the Stage 1 understanding-only path. Stage 2
sets `visual_gen=True` and uses the MoT `*_moe_gen` parameter siblings,
VAE latents prepared by the recipe, and the flow-matching MSE head.

**`backend`** `= resolve_bagel_backend(backend)`

---

**`backend_config_resolver`** `= staticmethod(resolve_bagel_backend)`

---

**`hidden_size`** `= config.text_config.hidden_size`

---

**`model`** `= BagelModel(config, backend=(self.backend))`

---

**`num_heads`** `= config.text_config.num_attention_heads`

---

**`state_dict_adapter`**

---

**`tie_word_embeddings_support`** `TieSupport = TieSupport.UNTIED_ONLY`

---

**`use_moe`**

---

```python
nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.forward(
    sequence_length: int,
    packed_text_ids: torch.LongTensor,
    packed_text_indexes: torch.LongTensor,
    sample_lens: typing.List[int],
    packed_position_ids: torch.LongTensor,
    nested_attention_masks: typing.List[torch.Tensor] | None = None,
    split_lens: typing.List[int] | None = None,
    attn_modes: typing.List[str] | None = None,
    packed_vit_tokens: torch.Tensor | None = None,
    packed_vit_token_indexes: torch.LongTensor | None = None,
    packed_vit_position_ids: torch.LongTensor | None = None,
    vit_token_seqlens: torch.Tensor | None = None,
    padded_latent: torch.Tensor | None = None,
    patchified_vae_latent_shapes: typing.List[typing.Tuple[int, int]] | None = None,
    packed_latent_position_ids: torch.LongTensor | None = None,
    packed_vae_token_indexes: torch.LongTensor | None = None,
    packed_timesteps: torch.Tensor | None = None,
    mse_loss_indexes: torch.Tensor | None = None,
    ce_loss_indexes: torch.Tensor | None = None,
    packed_label_ids: torch.Tensor | None = None,
    ce_loss_weights: torch.Tensor | None = None
) -> typing.Dict[str, torch.Tensor | None]
```

Run the BAGEL mixed-modal forward.

Stage 1 (`visual_gen=False`) skips the flow-matching branch and MSE
computation; Stage 2 activates both. `ce_loss_weights` is accepted
for data-pipeline compatibility but not consumed here - CE is returned
per-token (`reduction="none"`) and the trainer may apply weights
downstream.

Stage 2 inputs (`padded_latent`, `patchified_vae_latent_shapes`,
`packed_latent_position_ids`, `packed_vae_token_indexes`,
`packed_timesteps`, `mse_loss_indexes`) are produced by the
BAGEL collator when a pack contains t2i/edit samples. `padded_latent`
is the VAE-encoded latent tensor; recipe must call `vae_model.encode`
on the raw `padded_images` before forward; this module does not own
the VAE.

**Returns:** `Dict[str, torch.Tensor | None]`

`dict(ce=Tensor|None, mse=Tensor|None)` - both can be `None`

```python
nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.from_pretrained(
    pretrained_model_name_or_path: typing.Union[str, os.PathLike],
    stage: typing.Union[int, str] = 1,
    strict: bool = False,
    kwargs: typing.Any = {}
) -> 'BagelForUnifiedMultimodal'
```

classmethod

Load a BAGEL-7B-MoT checkpoint directory into this class.

Reads `config.json` via `BagelConfig.from_pretrained`, constructs
an empty model, and then loads `ema.safetensors` filtered by
`BagelStateDictAdapter`. Stage 2 VAE weights are loaded by the
recipe because the VAE is not owned by this module tree.

**Parameters:**

**`pretrained_model_name_or_path`** `Union[str, os.PathLike]`

Directory containing the HF-layout
BAGEL checkpoint.

---

**`stage`** `Union[int, str]` — default: 1

`1` (UND only) or `2` (UND + GEN). Strings
`"stage1"` / `"stage2"` are also accepted.

---

**`strict`** `bool` — default: False

If `True`, raise on state-dict keys that don't match the
adapter patterns. Defaults to `False` for compatibility with
checkpoint sidecar files.

---

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

Model overrides. `backend` is passed to the BAGEL model
constructor; the remaining values are forwarded to
`BagelConfig.from_pretrained`.

---

**Returns:** `'BagelForUnifiedMultimodal'`

A fully-initialized `BagelForUnifiedMultimodal` with weights

```python
nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.get_input_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.get_output_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.initialize_weights() -> None
```

Initialize BAGEL weights after AM materializes a `from_config` model.

`BagelForUnifiedMultimodal` is an `nn.Module` root, not a HF
`PreTrainedModel` root. AM's meta-device `from_config` path
materializes parameters after sharding and then calls this method.
Delegate Qwen/SigLIP subtrees to their HF-style initializers, then
initialize BAGEL-only connector and generation modules.

```python
nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.supports_config(
    config: typing.Any
) -> bool
```

classmethod

Return `True` if this custom class supports `config`.

```python
nemo_automodel.components.models.bagel.autoencoder.load_bagel_autoencoder(
    local_path: str | None
) -> tuple[nemo_automodel.components.models.bagel.autoencoder.AutoEncoder, nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams]
```

Load the BAGEL autoencoder from `ae.safetensors`.

**Parameters:**

**`local_path`** `str | None`

Local path to `ae.safetensors`. If `None`, the module is
returned with randomly initialized weights.

---

**Returns:** `tuple[AutoEncoder, AutoEncoderParams]`

The autoencoder module and its architecture parameters.