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

Qwen2 language backbone with packed-sequence attention and MoT shell.

Stage 1 uses `PackedAttention` + `Qwen2DecoderLayer`. The
`PackedAttentionMoT` / `Qwen2MoTDecoderLayer` shells are defined so that
the `*_moe_gen` parameter siblings exist in the module tree and survive
checkpoint round-tripping; they remain dormant in Stage 1 when
`packed_gen_token_indexes` is empty.

## Module Contents

### Classes

| Name                                                                                                               | Description                                                                       |
| ------------------------------------------------------------------------------------------------------------------ | --------------------------------------------------------------------------------- |
| [`BaseNavitOutputWithPast`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-BaseNavitOutputWithPast) | BAGEL packed decoder output with optional past key-value cache.                   |
| [`NaiveCache`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-NaiveCache)                           | Dict-backed KV cache, one entry per layer (BAGEL inference helper).               |
| [`PackedAttention`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-PackedAttention)                 | BAGEL's packed-sequence attention (UND path, no MoT).                             |
| [`PackedAttentionMoT`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-PackedAttentionMoT)           | MoT variant: adds `*_moe_gen` siblings of every projection and QK-norm.           |
| [`Qwen2DecoderLayer`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2DecoderLayer)             | Standard (non-MoT) packed Qwen2 decoder block.                                    |
| [`Qwen2ForCausalLM`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2ForCausalLM)               | Packed-sequence Qwen2 LM head wrapper.                                            |
| [`Qwen2MLP`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2MLP)                               | SwiGLU MLP used in Qwen2 (gate\_proj \* up\_proj -> down\_proj).                  |
| [`Qwen2MoTDecoderLayer`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2MoTDecoderLayer)       | MoT decoder: every norm/MLP is duplicated into `*_moe_gen` siblings.              |
| [`Qwen2Model`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2Model)                           | Packed-sequence Qwen2 backbone.                                                   |
| [`Qwen2PreTrainedModel`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2PreTrainedModel)       | Abstract base class — mirrors HF Qwen2PreTrainedModel flags.                      |
| [`Qwen2RMSNorm`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2RMSNorm)                       | Qwen2 RMSNorm (equivalent to T5LayerNorm).                                        |
| [`Qwen2RotaryEmbedding`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2RotaryEmbedding)       | Qwen2 rotary embedding — delegates inv\_freq init to HF `ROPE_INIT_FUNCTIONS`.    |
| [`_PackedAttentionBase`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_PackedAttentionBase)       | Common init for PackedAttention / PackedAttentionMoT (QKV shapes, RoPE, QK-norm). |

### Functions

| Name                                                                                                                                 | Description                                                                         |
| ------------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------- |
| [`_compute_default_rope_parameters`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_compute_default_rope_parameters) | Local "default" RoPE init — transformers 5.x dropped it from ROPE\_INIT\_FUNCTIONS. |
| [`_extract_rope_config`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_extract_rope_config)                         | Return a dict with `rope_theta` / scaling info, handling transformers 4.x and 5.x.  |
| [`_flash_attn_varlen`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_flash_attn_varlen)                             | -                                                                                   |
| [`_pad_sequence`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_pad_sequence)                                       | -                                                                                   |
| [`apply_rotary_pos_emb`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-apply_rotary_pos_emb)                         | Apply Rotary Position Embedding to query and key tensors.                           |
| [`rotate_half`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-rotate_half)                                           | Rotates half the hidden dims of the input.                                          |

### Data

[`_DECODER_LAYER_DICT`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_DECODER_LAYER_DICT)

[`__all__`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-__all__)

[`_flex_attention`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_flex_attention)

### API

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.BaseNavitOutputWithPast(
    packed_query_sequence: torch.FloatTensor = None,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None
)
```

Dataclass

**Bases:** `ModelOutput`

BAGEL packed decoder output with optional past key-value cache.

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache(
    num_layers: int
)
```

Dict-backed KV cache, one entry per layer (BAGEL inference helper).

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttention(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None
)
```

**Bases:** [\_PackedAttentionBase](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_PackedAttentionBase)

BAGEL's packed-sequence attention (UND path, no MoT).

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttention.forward(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttention.forward_inference(
    packed_query_sequence: torch.Tensor,
    query_lens: torch.Tensor,
    packed_query_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    packed_query_indexes: torch.Tensor,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None,
    key_values_lens: typing.Optional[torch.Tensor] = None,
    packed_key_value_indexes: typing.Optional[torch.Tensor] = None,
    update_past_key_values: bool = True,
    is_causal: bool = True
) -> typing.Tuple[torch.Tensor, typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache]]
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttention.forward_train(
    packed_sequence: torch.Tensor,
    sample_lens: typing.List[int],
    attention_mask,
    packed_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttentionMoT(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None
)
```

**Bases:** [\_PackedAttentionBase](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_PackedAttentionBase)

MoT variant: adds `*_moe_gen` siblings of every projection and QK-norm.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttentionMoT.forward(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttentionMoT.forward_inference(
    packed_query_sequence: torch.Tensor,
    query_lens: torch.Tensor,
    packed_query_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    packed_query_indexes: torch.Tensor,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None,
    key_values_lens: typing.Optional[torch.Tensor] = None,
    packed_key_value_indexes: typing.Optional[torch.Tensor] = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: typing.Optional[torch.Tensor] = None,
    packed_text_indexes: typing.Optional[torch.Tensor] = None
) -> typing.Tuple[torch.Tensor, typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache]]
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttentionMoT.forward_train(
    packed_sequence: torch.Tensor,
    sample_lens: typing.List[int],
    attention_mask,
    packed_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    packed_und_token_indexes: torch.LongTensor,
    packed_gen_token_indexes: torch.LongTensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2DecoderLayer(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None
)
```

**Bases:** `Module`

Standard (non-MoT) packed Qwen2 decoder block.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2DecoderLayer.forward(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2DecoderLayer.forward_inference(
    packed_query_sequence: torch.Tensor,
    query_lens: torch.Tensor,
    packed_query_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    packed_query_indexes: torch.Tensor,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None,
    key_values_lens: typing.Optional[torch.Tensor] = None,
    packed_key_value_indexes: typing.Optional[torch.Tensor] = None,
    update_past_key_values: bool = True,
    is_causal: bool = True
) -> typing.Tuple[torch.Tensor, typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache]]
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2DecoderLayer.forward_train(
    packed_sequence: torch.Tensor,
    sample_lens: typing.List[int],
    attention_mask,
    packed_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor]
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM(
    config: transformers.Qwen2Config
)
```

**Bases:** [Qwen2PreTrainedModel](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2PreTrainedModel)

Packed-sequence Qwen2 LM head wrapper.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.forward(
    args = (),
    kwargs = {}
) -> typing.Union[torch.Tensor, nemo_automodel.components.models.bagel.modeling_qwen2_packed.BaseNavitOutputWithPast]
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.forward_inference(
    packed_query_sequence: torch.Tensor,
    query_lens: torch.Tensor,
    packed_query_position_ids: torch.Tensor,
    packed_query_indexes: torch.Tensor,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None,
    key_values_lens: typing.Optional[torch.Tensor] = None,
    packed_key_value_indexes: typing.Optional[torch.Tensor] = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: typing.Optional[torch.Tensor] = None,
    packed_text_indexes: typing.Optional[torch.Tensor] = None
) -> nemo_automodel.components.models.bagel.modeling_qwen2_packed.BaseNavitOutputWithPast
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.forward_train(
    packed_sequence: torch.Tensor,
    sample_lens: typing.List[int],
    attention_mask,
    packed_position_ids: torch.Tensor,
    packed_und_token_indexes: typing.Optional[torch.LongTensor] = None,
    packed_gen_token_indexes: typing.Optional[torch.LongTensor] = None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.get_decoder() -> torch.nn.Module
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.get_input_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.get_output_embeddings() -> torch.nn.Module
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.init_moe() -> None
```

Seed `*_moe_gen` parameters from their UND siblings (Stage 1 cold-start).

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.set_decoder(
    decoder: torch.nn.Module
) -> None
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.set_input_embeddings(
    value: torch.nn.Module
) -> None
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM.set_output_embeddings(
    new_embeddings: torch.nn.Module
) -> None
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MLP(
    config: transformers.Qwen2Config
)
```

**Bases:** `Module`

SwiGLU MLP used in Qwen2 (gate\_proj \* up\_proj -> down\_proj).

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MLP.forward(
    hidden_state: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MoTDecoderLayer(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None,
    attn_module: type = PackedAttentionMoT
)
```

**Bases:** `Module`

MoT decoder: every norm/MLP is duplicated into `*_moe_gen` siblings.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MoTDecoderLayer.forward(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MoTDecoderLayer.forward_inference(
    packed_query_sequence: torch.Tensor,
    query_lens: torch.Tensor,
    packed_query_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    packed_query_indexes: torch.Tensor,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None,
    key_values_lens: typing.Optional[torch.Tensor] = None,
    packed_key_value_indexes: typing.Optional[torch.Tensor] = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: typing.Optional[torch.Tensor] = None,
    packed_text_indexes: typing.Optional[torch.Tensor] = None
) -> typing.Tuple[torch.Tensor, typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache]]
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MoTDecoderLayer.forward_train(
    packed_sequence: torch.Tensor,
    sample_lens: typing.List[int],
    attention_mask,
    packed_position_embeddings: typing.Tuple[torch.Tensor, torch.Tensor],
    packed_und_token_indexes: torch.LongTensor,
    packed_gen_token_indexes: torch.LongTensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model(
    config: transformers.Qwen2Config
)
```

**Bases:** [Qwen2PreTrainedModel](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-Qwen2PreTrainedModel)

Packed-sequence Qwen2 backbone.

Selects `Qwen2DecoderLayer` or `Qwen2MoTDecoderLayer` per-layer based
on `config.layer_module` (string -> class). When the MoT variant is
active, `self.use_moe == True` and an extra `norm_moe_gen` sibling is
created for the final RMSNorm.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model.forward(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model.forward_inference(
    packed_query_sequence: torch.Tensor,
    query_lens: torch.Tensor,
    packed_query_position_ids: torch.Tensor,
    packed_query_indexes: torch.Tensor,
    past_key_values: typing.Optional[nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache] = None,
    key_values_lens: typing.Optional[torch.Tensor] = None,
    packed_key_value_indexes: typing.Optional[torch.Tensor] = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: typing.Optional[torch.Tensor] = None,
    packed_text_indexes: typing.Optional[torch.Tensor] = None
) -> nemo_automodel.components.models.bagel.modeling_qwen2_packed.BaseNavitOutputWithPast
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model.forward_train(
    packed_sequence: torch.Tensor,
    sample_lens: typing.List[int],
    attention_mask,
    packed_position_ids: torch.Tensor,
    packed_und_token_indexes: typing.Optional[torch.LongTensor] = None,
    packed_gen_token_indexes: typing.Optional[torch.LongTensor] = None
) -> torch.Tensor
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model.init_moe() -> None
```

Copy UND weights into MoE-gen siblings (Stage 1 cold-start seeding).

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2PreTrainedModel()
```

**Bases:** `PreTrainedModel`

Abstract base class — mirrors HF Qwen2PreTrainedModel flags.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2PreTrainedModel._init_weights(
    module: torch.nn.Module
) -> None
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2RMSNorm(
    hidden_size: int,
    eps: float = 1e-06
)
```

**Bases:** `Module`

Qwen2 RMSNorm (equivalent to T5LayerNorm).

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2RMSNorm.extra_repr() -> str
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2RMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2RotaryEmbedding(
    config: transformers.Qwen2Config,
    device: typing.Optional[torch.device] = None
)
```

**Bases:** `Module`

Qwen2 rotary embedding — delegates inv\_freq init to HF `ROPE_INIT_FUNCTIONS`.

DIVERGENCE: transformers 5.x removed `"default"` from `ROPE_INIT_FUNCTIONS`
so we fall back to a local copy of the pre-5.x default implementation when
the rope\_type is unspecified.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2RotaryEmbedding._dynamic_frequency_update(
    position_ids: torch.Tensor,
    device: torch.device
) -> None
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2RotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> typing.Tuple[torch.Tensor, torch.Tensor]
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed._PackedAttentionBase(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None
)
```

**Bases:** `Module`

Common init for PackedAttention / PackedAttentionMoT (QKV shapes, RoPE, QK-norm).

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._compute_default_rope_parameters(
    config: transformers.Qwen2Config,
    device: typing.Optional[torch.device] = None,
    seq_len: typing.Optional[int] = None
) -> typing.Tuple[torch.Tensor, float]
```

Local "default" RoPE init — transformers 5.x dropped it from ROPE\_INIT\_FUNCTIONS.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._extract_rope_config(
    config: transformers.Qwen2Config
) -> dict
```

Return a dict with `rope_theta` / scaling info, handling transformers 4.x and 5.x.

DIVERGENCE: upstream BAGEL was written against transformers 4.4x where
`Qwen2Config` exposes `rope_theta` and `rope_scaling` as top-level
attributes. transformers 5.x moves these into a single `rope_parameters`
dict on Qwen2Config (Llama still keeps the old layout). AM's container runs
transformers 5.x, so we normalize here instead of hard-coding one schema.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._flash_attn_varlen(
    args = (),
    kwargs = {}
)
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._pad_sequence(
    tensor: torch.Tensor,
    pad_size: int
) -> torch.Tensor
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    position_ids: typing.Optional[torch.Tensor] = None,
    unsqueeze_dim: int = 1
) -> typing.Tuple[torch.Tensor, torch.Tensor]
```

Apply Rotary Position Embedding to query and key tensors.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.rotate_half(
    x: torch.Tensor
) -> torch.Tensor
```

Rotates half the hidden dims of the input.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._DECODER_LAYER_DICT = {'Qwen2DecoderLayer': Qwen2DecoderLayer, 'Qwen2MoTDecoderLayer': partial(Qwen2Mo...
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.__all__ = ['Qwen2RMSNorm', 'Qwen2RotaryEmbedding', 'Qwen2MLP', 'PackedAttention', 'PackedA...
```

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._flex_attention = FlexAttention.flex_attn
```