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`	BAGEL packed decoder output with optional past key-value cache.
`NaiveCache`	Dict-backed KV cache, one entry per layer (BAGEL inference helper).
`PackedAttention`	BAGEL’s packed-sequence attention (UND path, no MoT).
`PackedAttentionMoT`	MoT variant: adds `*_moe_gen` siblings of every projection and QK-norm.
`Qwen2DecoderLayer`	Standard (non-MoT) packed Qwen2 decoder block.
`Qwen2ForCausalLM`	Packed-sequence Qwen2 LM head wrapper.
`Qwen2MLP`	SwiGLU MLP with an independently configurable linear backend.
`Qwen2MoTDecoderLayer`	MoT decoder: every norm/MLP is duplicated into `*_moe_gen` siblings.
`Qwen2Model`	Packed-sequence Qwen2 backbone.
`Qwen2PreTrainedModel`	Abstract base class — mirrors HF Qwen2PreTrainedModel flags.
`Qwen2RMSNorm`	Qwen2 RMSNorm (equivalent to T5LayerNorm).
`Qwen2RotaryEmbedding`	Qwen2 rotary embedding — delegates inv_freq init to HF `ROPE_INIT_FUNCTIONS`.
`_PackedAttentionBase`	Common init for PackedAttention / PackedAttentionMoT (QKV shapes, RoPE, QK-norm).

Functions

Name	Description
`_apply_qk_norm`	Apply Q/K RMSNorm over each attention head’s final dimension.
`_bagel_fused_rope`	Fused RoPE apply: rotate `q` and `k` in a single compiled pointwise kernel.
`_bagel_fused_silu_mul`	Fused SwiGLU gate: `silu(gate) * up` in a single compiled pointwise kernel.
`_compute_default_rope_parameters`	Local “default” RoPE init — transformers 5.x dropped it from ROPE_INIT_FUNCTIONS.
`_extract_rope_config`	Return a dict with `rope_theta` / scaling info, handling transformers 4.x and 5.x.
`_flash_attn_varlen`	-
`_index_put_matching_dtype`	Assign selected packed-token rows and surface backend dtype changes once.
`_initialize_linear`	Construct a torch-compatible linear while preserving the configured parameter dtype.
`_initialize_rms_norm`	Construct the selected RMSNorm while preserving BAGEL’s torch baseline.
`_pad_sequence`	-
`apply_rotary_pos_emb`	Apply Rotary Position Embedding to query and key tensors.
`rotate_half`	Rotates half the hidden dims of the input.

Data

_DECODER_LAYER_DICT

_WARNED_INDEX_PUT_DTYPE_CASTS

__all__

_flex_attention

logger

API

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.

packed_query_sequence

FloatTensor = None

past_key_values

Optional[NaiveCache] = None

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, Optional[Tensor]] = {k: None for k in (range(num_layers))}

num_layers

int

seq_lens

int

value_cache

dict[int, Optional[Tensor]] = {k: None for k in (range(num_layers))}

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

Bases: _PackedAttentionBase

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

k_norm

q_norm

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

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

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

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

Bases: _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

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

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

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: typing.Optional[torch.LongTensor] = None,
    mot_inv: typing.Optional[torch.LongTensor] = None
) -> torch.Tensor

class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2DecoderLayer(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None,
    backend: typing.Optional[nemo_automodel.components.models.bagel.configuration.BagelBackendConfig] = 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

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

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

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

class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2ForCausalLM(
    config: transformers.Qwen2Config,
    backend: typing.Optional[nemo_automodel.components.models.bagel.configuration.BagelBackendConfig] = None
)

Bases: 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

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]

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

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,
    mot_perm: typing.Optional[torch.LongTensor] = None,
    mot_inv: typing.Optional[torch.LongTensor] = None
) -> torch.Tensor

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

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

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

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

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

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

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

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

class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MLP(
    config: transformers.Qwen2Config,
    backend: typing.Optional[nemo_automodel.components.models.bagel.configuration.BagelBackendConfig] = 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

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

class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2MoTDecoderLayer(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None,
    attn_module: type = PackedAttentionMoT,
    backend: typing.Optional[nemo_automodel.components.models.bagel.configuration.BagelBackendConfig] = 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

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

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

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: typing.Optional[torch.LongTensor] = None,
    mot_inv: typing.Optional[torch.LongTensor] = None
) -> torch.Tensor

class nemo_automodel.components.models.bagel.modeling_qwen2_packed.Qwen2Model(
    config: transformers.Qwen2Config,
    backend: typing.Optional[nemo_automodel.components.models.bagel.configuration.BagelBackendConfig] = None
)

Bases: 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

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

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

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,
    mot_perm: typing.Optional[torch.LongTensor] = None,
    mot_inv: typing.Optional[torch.LongTensor] = None
) -> torch.Tensor

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

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

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'

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

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

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

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

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.

max_seq_len_cached

= config.max_position_embeddings

original_inv_freq

= self.inv_freq

original_max_seq_len

= config.max_position_embeddings

rope_type

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

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

class nemo_automodel.components.models.bagel.modeling_qwen2_packed._PackedAttentionBase(
    config: transformers.Qwen2Config,
    layer_idx: typing.Optional[int] = None,
    backend: typing.Optional[nemo_automodel.components.models.bagel.configuration.BagelBackendConfig] = 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

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

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:

torch.Tensor

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

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.

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.

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.

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.

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

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.

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.

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.

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

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,
    fused: bool = False
) -> typing.Tuple[torch.Tensor, torch.Tensor]

Apply Rotary Position Embedding to query and key tensors.

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

Rotates half the hidden dims of the input.

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

nemo_automodel.components.models.bagel.modeling_qwen2_packed._WARNED_INDEX_PUT_DTYPE_CASTS: set[tuple[dtype, dtype]] = set()

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

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

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