> 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 with an independently configurable linear backend.                     |
| [`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                                                                          |
| ------------------------------------------------------------------------------------------------------------------------------------ | ------------------------------------------------------------------------------------ |
| [`_apply_qk_norm`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_apply_qk_norm)                                     | Apply Q/K RMSNorm over each attention head's final dimension.                        |
| [`_bagel_fused_rope`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_bagel_fused_rope)                               | Fused RoPE apply: rotate `q` and `k` in a single compiled pointwise kernel.          |
| [`_bagel_fused_silu_mul`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_bagel_fused_silu_mul)                       | Fused SwiGLU gate: `silu(gate) * up` in a single compiled pointwise kernel.          |
| [`_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)                             | -                                                                                    |
| [`_index_put_matching_dtype`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_index_put_matching_dtype)               | Assign selected packed-token rows and surface backend dtype changes once.            |
| [`_initialize_linear`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_initialize_linear)                             | Construct a torch-compatible linear while preserving the configured parameter dtype. |
| [`_initialize_rms_norm`](#nemo_automodel-components-models-bagel-modeling_qwen2_packed-_initialize_rms_norm)                         | Construct the selected RMSNorm while preserving BAGEL's torch baseline.              |
| [`_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)

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

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

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

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

### API

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

Dataclass

**Bases:** `ModelOutput`

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

**`packed_query_sequence`** `FloatTensor = None`

---

**`past_key_values`** `NaiveCache | None = None`

---

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

**`key_cache`** `dict[int, Tensor | None] = {k: None for k in (range(num_layers))}`

---

**`num_layers`** `int`

---

**`seq_lens`** `int`

---

**`value_cache`** `dict[int, Tensor | None] = {k: None for k in (range(num_layers))}`

---

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.PackedAttention(
    config: transformers.Qwen2Config,
    layer_idx: int | None = None,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

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

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

**`k_norm`**

---

**`q_norm`**

---

```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: nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None = None,
    key_values_lens: torch.Tensor | None = None,
    packed_key_value_indexes: torch.Tensor | None = None,
    update_past_key_values: bool = True,
    is_causal: bool = True
) -> typing.Tuple[torch.Tensor, nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None]
```

```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: int | None = None,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

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

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

**`k_norm`**

---

**`k_norm_moe_gen`**

---

**`k_proj_moe_gen`**

---

**`o_proj_moe_gen`**

---

**`q_norm`**

---

**`q_norm_moe_gen`**

---

**`q_proj_moe_gen`**

---

**`v_proj_moe_gen`**

---

```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: nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None = None,
    key_values_lens: torch.Tensor | None = None,
    packed_key_value_indexes: torch.Tensor | None = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: torch.Tensor | None = None,
    packed_text_indexes: torch.Tensor | None = None
) -> typing.Tuple[torch.Tensor, nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None]
```

```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,
    mot_perm: torch.LongTensor | None = None,
    mot_inv: torch.LongTensor | None = None
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2DecoderLayer(
    config: transformers.Qwen2Config,
    layer_idx: int | None = None,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

**Bases:** `Module`

Standard (non-MoT) packed Qwen2 decoder block.

**`backend`** `= backend or BagelBackendConfig()`

---

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

---

**`input_layernorm`**

---

**`mlp`** `= Qwen2MLP(config, backend=(self.backend))`

---

**`post_attention_layernorm`**

---

**`self_attn`**

---

```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: nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None = None,
    key_values_lens: torch.Tensor | None = None,
    packed_key_value_indexes: torch.Tensor | None = None,
    update_past_key_values: bool = True,
    is_causal: bool = True
) -> typing.Tuple[torch.Tensor, nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None]
```

```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,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

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

Packed-sequence Qwen2 LM head wrapper.

**`_tied_weights_keys`** `= ['lm_head.weight']`

---

**`backend`** `= backend or BagelBackendConfig()`

---

**`lm_head`**

---

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

---

**`vocab_size`** `= config.vocab_size`

---

```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: nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None = None,
    key_values_lens: torch.Tensor | None = None,
    packed_key_value_indexes: torch.Tensor | None = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: torch.Tensor | None = None,
    packed_text_indexes: torch.Tensor | None = 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: torch.LongTensor | None = None,
    packed_gen_token_indexes: torch.LongTensor | None = None,
    mot_perm: torch.LongTensor | None = None,
    mot_inv: torch.LongTensor | None = 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,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

**Bases:** `Module`

SwiGLU MLP with an independently configurable linear backend.

**`_fuse_silu_mul`**

---

**`act_fn`** `= ACT2FN[config.hidden_act]`

---

**`backend`** `= backend or BagelBackendConfig()`

---

**`down_proj`**

---

**`gate_proj`**

---

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

---

**`intermediate_size`** `= config.intermediate_size`

---

**`up_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: int | None = None,
    attn_module: type = PackedAttentionMoT,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

**Bases:** `Module`

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

**`backend`** `= backend or BagelBackendConfig()`

---

**`freeze_und`** `= getattr(config, 'freeze_und', False)`

---

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

---

**`input_layernorm`**

---

**`input_layernorm_moe_gen`**

---

**`mlp`** `= Qwen2MLP(config, backend=(self.backend))`

---

**`mlp_moe_gen`** `= Qwen2MLP(config, backend=(self.backend))`

---

**`post_attention_layernorm`**

---

**`post_attention_layernorm_moe_gen`**

---

**`self_attn`**

---

```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: nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None = None,
    key_values_lens: torch.Tensor | None = None,
    packed_key_value_indexes: torch.Tensor | None = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: torch.Tensor | None = None,
    packed_text_indexes: torch.Tensor | None = None
) -> typing.Tuple[torch.Tensor, nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None]
```

```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,
    mot_perm: torch.LongTensor | None = None,
    mot_inv: torch.LongTensor | None = None
) -> torch.Tensor
```

```python
class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model(
    config: transformers.Qwen2Config,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

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

**`backend`** `= backend or BagelBackendConfig()`

---

**`embed_tokens`**

---

**`layers`**

---

**`norm`**

---

**`norm_moe_gen`**

---

**`padding_idx`** `= config.pad_token_id`

---

**`rotary_emb`** `= Qwen2RotaryEmbedding(config=config)`

---

**`use_moe`** `= 'Mo' in layer_module_name`

---

**`vocab_size`** `= config.vocab_size`

---

```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: nemo_automodel.components.models.bagel.modeling_qwen2_packed.NaiveCache | None = None,
    key_values_lens: torch.Tensor | None = None,
    packed_key_value_indexes: torch.Tensor | None = None,
    update_past_key_values: bool = True,
    is_causal: bool = True,
    mode: str = 'und',
    packed_vae_token_indexes: torch.Tensor | None = None,
    packed_text_indexes: torch.Tensor | None = 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: torch.LongTensor | None = None,
    packed_gen_token_indexes: torch.LongTensor | None = None,
    mot_perm: torch.LongTensor | None = None,
    mot_inv: torch.LongTensor | None = 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.

**`_no_split_modules`** `= ['Qwen2DecoderLayer', 'Qwen2MoTDecoderLayer']`

---

**`_skip_keys_device_placement`** `= 'past_key_values'`

---

**`base_model_prefix`** `= 'model'`

---

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

**`weight`** `= nn.Parameter(torch.ones(hidden_size))`

---

```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: torch.device | None = 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.

**`max_seq_len_cached`** `= config.max_position_embeddings`

---

**`original_inv_freq`** `= self.inv_freq`

---

**`original_max_seq_len`** `= config.max_position_embeddings`

---

**`rope_type`**

---

```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: int | None = None,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)
```

**Bases:** `Module`

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

**`attention_dropout`** `= config.attention_dropout`

---

**`backend`** `= backend or BagelBackendConfig()`

---

**`head_dim`** `= self.hidden_size // self.num_heads`

---

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

---

**`is_causal`** `= getattr(config, 'is_causal', True)`

---

**`k_proj`**

---

**`max_position_embeddings`** `= config.max_position_embeddings`

---

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

---

**`num_key_value_groups`** `= self.num_heads // self.num_key_value_heads`

---

**`num_key_value_heads`** `= config.num_key_value_heads`

---

**`o_proj`**

---

**`q_proj`**

---

**`rope_theta`**

---

**`v_proj`**

---

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._apply_qk_norm(
    norm: torch.nn.Module,
    hidden_states: torch.Tensor,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig,
    eps: float
) -> torch.Tensor
```

Apply Q/K RMSNorm over each attention head's final dimension.

**Parameters:**

**`norm`** `nn.Module`

RMSNorm module to apply.

---

**`hidden_states`** `torch.Tensor`

Tensor of shape `[tokens, heads, head_dim]`. RMSNorm
is computed over the final `head_dim` axis.

---

**`backend`** `BagelBackendConfig`

Resolved BAGEL backend configuration.

---

**`eps`** `float`

Epsilon used by the explicit FP32 Transformer Engine path.

---

**Returns:** `torch.Tensor`

Tensor of shape `[tokens, heads, head_dim]` with the same layout as

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._bagel_fused_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> typing.Tuple[torch.Tensor, torch.Tensor]
```

Fused RoPE apply: rotate `q` and `k` in a single compiled pointwise kernel.

`cos`/`sin` arrive already unsqueezed to broadcast over the head dim. Compiling this
pointwise region — like the SwiGLU gate, over regular activation tensors and no FSDP2
DTensor parameters — fuses the `rotate_half` slice/negate/cat and the mul/add into far
fewer kernels than eager ATen, with no recompile thrash.

**Parameters:**

**`q`** `torch.Tensor`

Query states, shape `[tokens, heads, head_dim]`.

---

**`k`** `torch.Tensor`

Key states, shape `[tokens, kv_heads, head_dim]`.

---

**`cos`** `torch.Tensor`

Cosine table, broadcastable to `q` and `k`.

---

**`sin`** `torch.Tensor`

Sine table, broadcastable to `q` and `k`.

---

**Returns:** `Tuple[torch.Tensor, torch.Tensor]`

Tuple of rotated `(q_embed, k_embed)` matching the inputs' shapes and dtypes.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._bagel_fused_silu_mul(
    gate: torch.Tensor,
    up: torch.Tensor
) -> torch.Tensor
```

Fused SwiGLU gate: `silu(gate) * up` in a single compiled pointwise kernel.

Compiling ONLY this pointwise activation — over regular (non-DTensor) activation tensors,
not the projections/attention that run on FSDP2 DTensor params — fuses the silu and the
multiply into one kernel (fewer launches, no materialization of the silu output) without
triggering the DTensor-spec recompile thrash that whole-layer `torch.compile` hits.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._compute_default_rope_parameters(
    config: transformers.Qwen2Config,
    device: torch.device | None = None,
    seq_len: int | None = 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._index_put_matching_dtype(
    destination: torch.Tensor,
    index: torch.Tensor,
    source: torch.Tensor
) -> None
```

Assign selected packed-token rows and surface backend dtype changes once.

**Parameters:**

**`destination`** `torch.Tensor`

Packed tensor of shape `[tokens, ...]`. This tensor is
mutated in place.

---

**`index`** `torch.Tensor`

One-dimensional integer tensor of shape `[selected_tokens]`
selecting rows along the packed-token axis.

---

**`source`** `torch.Tensor`

Tensor of shape `[selected_tokens, ...]` whose trailing
dimensions match `destination`.

---

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._initialize_linear(
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig,
    in_features: int,
    out_features: int,
    bias: bool
) -> torch.nn.Module
```

Construct a torch-compatible linear while preserving the configured parameter dtype.

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed._initialize_rms_norm(
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig,
    hidden_size: int,
    eps: float
) -> torch.nn.Module
```

Construct the selected RMSNorm while preserving BAGEL's torch baseline.

```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: torch.Tensor | None = None,
    unsqueeze_dim: int = 1,
    fused: bool = False
) -> 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._WARNED_INDEX_PUT_DTYPE_CASTS: set[tuple[dtype, dtype]] = set()
```

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

```python
nemo_automodel.components.models.bagel.modeling_qwen2_packed.logger = logging.getLogger(__name__)
```