nemo_automodel.components.models.qwen3_omni_moe.model

Module Contents

Classes

Name	Description
`Qwen3OmniMoeThinkerForConditionalGeneration`	Qwen3OmniMoe Thinker for Conditional Generation with multimodal support.
`Qwen3OmniMoeThinkerTextModel`	Qwen3OmniMoe Thinker Text Model with MRoPE and sparse MoE layers.
`Qwen3OmniMoeThinkerTextRotaryEmbedding`	Ensure HF rotary frequency buffers stay fp32 across module-wide casts.

Data

ModelClass

API

class nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration(
    config: transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe.Qwen3OmniMoeThinkerConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

Bases: HFCheckpointingMixin, HFQwen3OmniMoeThinkerForConditionalGeneration, MoEFSDPSyncMixin

Qwen3OmniMoe Thinker for Conditional Generation with multimodal support.

lm_head

model

num_experts

= text_config.num_experts

num_experts_per_tok

= text_config.num_experts_per_tok

pad_token_id

router_aux_loss_coef

= getattr(text_config, 'router_aux_loss_coef', 0.0)

spatial_merge_size

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

vocab_size

= text_config.vocab_size

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.forward(
    input_ids: torch.Tensor,
    input_features: torch.FloatTensor | None = None,
    pixel_values: torch.FloatTensor | None = None,
    pixel_values_videos: torch.FloatTensor | None = None,
    image_grid_thw: torch.LongTensor | None = None,
    video_grid_thw: torch.LongTensor | None = None,
    attention_mask: torch.Tensor | None = None,
    feature_attention_mask: torch.LongTensor | None = None,
    audio_feature_lengths: torch.LongTensor | None = None,
    position_ids: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    inputs_embeds: torch.FloatTensor | None = None,
    labels: torch.LongTensor | None = None,
    output_router_logits: bool | None = None,
    use_audio_in_video: bool | None = None,
    video_second_per_grid: torch.Tensor | None = None,
    logits_to_keep: typing.Union[int, torch.Tensor] = 0,
    output_hidden_states: bool | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor | dict | transformers.modeling_outputs.CausalLMOutputWithPast

Forward pass with multimodal fusion.

Parameters:

input_ids

torch.Tensor

Integer token IDs of shape [batch, sequence].

input_features

torch.FloatTensor | NoneDefaults to None

Audio features of shape [audios, mel_bins, frames].

pixel_values

torch.FloatTensor | NoneDefaults to None

Packed image patches of shape [image_patches, patch_features], where patch_features flattens channels, temporal depth, height, and width.

pixel_values_videos

torch.FloatTensor | NoneDefaults to None

Packed video patches of shape [video_patches, patch_features].

image_grid_thw

torch.LongTensor | NoneDefaults to None

Integer temporal/height/width grids of shape [images, 3].

video_grid_thw

torch.LongTensor | NoneDefaults to None

Integer temporal/height/width grids of shape [videos, 3].

attention_mask

torch.Tensor | NoneDefaults to None

Padding mask of shape [batch, sequence].

feature_attention_mask

torch.LongTensor | NoneDefaults to None

Valid audio frames of shape [audios, frames].

audio_feature_lengths

torch.LongTensor | NoneDefaults to None

Valid frame counts of shape [audios].

position_ids

torch.Tensor | NoneDefaults to None

Rotary coordinates of shape [3, batch, sequence].

padding_mask

torch.Tensor | NoneDefaults to None

Valid token mask of shape [batch, sequence].

inputs_embeds

torch.FloatTensor | NoneDefaults to None

Optional embeddings of shape [batch, sequence, hidden].

labels

torch.LongTensor | NoneDefaults to None

Target token IDs of shape [batch, sequence].

output_router_logits

bool | NoneDefaults to None

Whether to output router logits

use_audio_in_video

bool | NoneDefaults to None

Whether audio is in video

video_second_per_grid

torch.Tensor | NoneDefaults to None

Temporal grid spacing of shape [videos].

logits_to_keep

Union[int, torch.Tensor]Defaults to 0

Integer count or index tensor of shape [selected_positions]. If > 0, only compute logits for the last logits_to_keep token positions (0 = all positions). Enables memory-efficient fused cross-entropy by letting the recipe request a single-position lm_head projection alongside the final hidden states.

output_hidden_states

bool | NoneDefaults to None

When set, the returned output carries the final hidden states (the input to lm_head) so the recipe can run fused linear cross-entropy.

**attn_kwargs

AnyDefaults to {}

Additional attention arguments

Returns: torch.Tensor | dict | CausalLMOutputWithPast

Logits of shape [batch, selected_sequence, vocab] (or [tokens, vocab]

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.from_config(
    config: transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe.Qwen3OmniMoeThinkerConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    backend: nemo_automodel.components.models.common.BackendConfig | None = None,
    kwargs = {}
)

classmethod

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.from_pretrained(
    pretrained_model_name_or_path: str,
    model_args = (),
    kwargs = {}
)

classmethod

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.get_input_embeddings()

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.get_output_embeddings()

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.initialize_weights(
    buffer_device: torch.device | None = None,
    dtype: torch.dtype = torch.bfloat16
) -> None

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.set_input_embeddings(
    value
)

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerForConditionalGeneration.set_output_embeddings(
    new_embeddings
)

class nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextModel(
    config: transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe.Qwen3OmniMoeTextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None,
    moe_overrides: dict | None = None
)

Bases: Module

Qwen3OmniMoe Thinker Text Model with MRoPE and sparse MoE layers.

embed_tokens

layers

moe_config

= moe_config or MoEConfig(**moe_defaults)

norm

padding_idx

= getattr(config, 'pad_token_id', None)

rotary_emb

= Qwen3OmniMoeThinkerTextRotaryEmbedding(config)

vocab_size

= config.vocab_size

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextModel._deepstack_process(
    hidden_states,
    visual_pos_masks,
    visual_embeds
)

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextModel.forward(
    input_ids: torch.Tensor | None = None,
    inputs_embeds: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    visual_pos_masks: torch.Tensor | None = None,
    deepstack_visual_embeds: list[torch.Tensor] | None = None,
    attn_kwargs: typing.Any = {}
) -> torch.Tensor

visual_pos_masks (torch.Tensor of shape (batch_size, seqlen), optional): The mask of the visual positions. deepstack_visual_embeds (list[torch.Tensor], optional): The deepstack visual embeddings. The shape is (num_layers, visual_seqlen, embed_dim). The feature is extracted from the different visual encoder layers, and fed to the decoder hidden states. It’s from the paper DeepStack(https://arxiv.org/abs/2406.04334).

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextModel.init_weights(
    buffer_device: torch.device | None = None
) -> None

class nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextRotaryEmbedding()

Bases: HFQwen3OmniMoeThinkerTextRotaryEmbedding

Ensure HF rotary frequency buffers stay fp32 across module-wide casts.

nemo_automodel.components.models.qwen3_omni_moe.model.Qwen3OmniMoeThinkerTextRotaryEmbedding._apply(
    fn: typing.Any,
    recurse: bool = True
)

nemo_automodel.components.models.qwen3_omni_moe.model.ModelClass = Qwen3OmniMoeThinkerForConditionalGeneration