bridge.models.nemotron_omni.modeling_nemotron_omni

Nemotron Omni model for processor-expanded multimodal token sequences.

Unlike MCore LLaVAModel, it does not collapse a run of image placeholders before the model and reconstruct it inside the model. The processor-provided sequence already contains one image placeholder per projected RADIO feature. This model replaces those positions in place. Packing is completed by the collator; the model only applies context-parallel sharding after media insertion.

The historical collapse/expand implementation remains available explicitly as NemotronOmniLlavaModel for compatibility with existing checkpoints, but it is not the canonical model selected by AutoBridge.

Module Contents

Classes

NemotronOmniModel

Nemotron Omni model whose input sequence is already media-expanded.

Functions

_ignore_transformer_engine_extra_state	Allow checkpoints produced before Transformer Engine added extra state.
_build_vision_packed_seq_params	Build RADIO’s per-image THD boundaries from image sizes.
_pixel_shuffle_dynamic_resolution	Group each spatial 2x2 patch block into the channel dimension.

API

bridge.models.nemotron_omni.modeling_nemotron_omni._ignore_transformer_engine_extra_state(

module: torch.nn.Module,

incompatible_keys: collections.namedtuple,

) → None

Allow checkpoints produced before Transformer Engine added extra state.

bridge.models.nemotron_omni.modeling_nemotron_omni._build_vision_packed_seq_params(

imgs_sizes: torch.Tensor,

patch_dim: int,

) → megatron.core.packed_seq_params.PackedSeqParams

Build RADIO’s per-image THD boundaries from image sizes.

bridge.models.nemotron_omni.modeling_nemotron_omni._pixel_shuffle_dynamic_resolution(

features: torch.Tensor,

*,

height: int,

width: int,

) → torch.Tensor

Group each spatial 2x2 patch block into the channel dimension.

A plain reshape groups four adjacent elements in the flattened sequence, which is not the same operation for a row-major non-square patch grid. Keep the spatial permutation identical to the historical Omni LLaVA path and the HF/vLLM implementation.

class bridge.models.nemotron_omni.modeling_nemotron_omni.NemotronOmniModel(

*,

language_transformer_config: megatron.core.transformer.transformer_config.TransformerConfig,

language_transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec,

language_vocab_size: int,

language_max_sequence_length: int,

vision_transformer_config: megatron.core.transformer.transformer_config.TransformerConfig,

vision_transformer_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec,

vision_projection_config: megatron.core.transformer.transformer_config.TransformerConfig,

vision_projection_layer_spec: megatron.core.transformer.spec_utils.ModuleSpec,

image_token_index: int,

parallel_output: bool = True,

share_embeddings_and_output_weights: bool = False,

language_position_embedding_type: str = 'rope',

pre_process: bool = True,

post_process: bool = True,

add_encoder: bool = True,

add_decoder: bool = True,

hybrid_layer_pattern: Optional[str] = None,

img_h: int = 512,

img_w: int = 512,

patch_dim: int = 16,

dynamic_resolution: bool = True,

vision_class_token_len: int = 10,

radio_force_eval_mode: bool = False,

radio_force_cpe_eval_mode: bool = False,

radio_interpolate_only_cpe: bool = False,

radio_cpe_aspect_ratio_select: bool = False,

radio_disable_cpe: bool = False,

temporal_patch_dim: int = 1,

separate_video_embedder: bool = False,

temporal_ckpt_compat: bool = False,

sound_model: Optional[torch.nn.Module] = None,

sound_projection: Optional[torch.nn.Module] = None,

sound_token_index: int = 0,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

vp_stage: Optional[int] = None,

)

Bases: megatron.core.transformer.MegatronModule

Nemotron Omni model whose input sequence is already media-expanded.

The collator supplies either a dense batch or a complete MCore THD stream. Media insertion is one-for-one and therefore length preserving. With context parallelism, the model inserts media into the full stream and then selects the rank-local CP shard without changing packed metadata.

Image, video, sound, and text inputs use the same one-feature-per-placeholder contract.

Initialization

model_owns_packing

False

model_owns_mtp_loss_mask_packing

False

model_slices_context_parallel_inputs

True

shared_embedding_or_output_weight()

Expose the language embedding for Megatron gradient finalization.

set_input_tensor(input_tensor) → None

Set the pipeline input on the language model.

freeze(

*,

freeze_language_model: bool = False,

freeze_vision_model: bool = False,

freeze_vision_projection: bool = False,

freeze_sound_model: bool = False,

freeze_sound_projection: bool = False,

) → None

Freeze selected leaf components.

static _merge_projected_media(

language_embeddings: torch.Tensor,

input_ids: torch.Tensor,

media_embeddings: torch.Tensor,

media_token_id: int,

attention_mask: Optional[torch.Tensor],

) → torch.Tensor

Replace each valid media placeholder with exactly one feature row.

attention_mask is a token-validity mask here, not MCore’s 4-D causal attention mask. Requiring an exact shape match prevents a causal mask from broadcasting the placeholder mask to [B, 1, S, S].

_encode_images(

images: torch.Tensor,

imgs_sizes: Optional[torch.Tensor],

vision_packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams],

num_frames: Optional[torch.Tensor],

) → torch.Tensor

Encode dynamic-resolution images and return one row per image token.

_encode_sound(

sound_clips: torch.Tensor,

sound_length: Optional[torch.Tensor],

) → torch.Tensor

Encode mel features and return valid projected rows in sample order.

_patchify_dynamic_images(

images: torch.Tensor,

imgs_sizes: torch.Tensor,

) → torch.Tensor

Convert padded processor pixels to RADIO’s packed patch representation.

The processor emits [num_images, channels, padded_height, padded_width]. RADIO’s dynamic-resolution path consumes [1, total_patches, channels * patch_dim**2]. Keeping this conversion here makes raw media tensors part of the model contract and avoids an Omni-only NeMo-RL pre-forward adapter. Already-patchified inputs remain accepted for Bridge/SFT callers.

static _select_sequence(

tensor: Optional[torch.Tensor],

index: torch.Tensor,

*,

dim: int,

) → Optional[torch.Tensor]

Select one CP shard from a token-aligned tensor.

_context_parallel_index(

*,

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams],

total_tokens: int,

device: torch.device,

) → Optional[torch.Tensor]

Build this rank’s CP index without changing sequence metadata.

_apply_context_parallel_sharding(

*,

input_ids: Optional[torch.Tensor],

combined_embeddings: Optional[torch.Tensor],

position_ids: Optional[torch.Tensor],

attention_mask: Optional[torch.Tensor],

labels: Optional[torch.Tensor],

loss_mask: Optional[torch.Tensor],

padding_mask: Optional[torch.Tensor],

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams],

) → tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], bool]

Apply one shared CP index after length-preserving media insertion.

forward(

input_ids: Optional[torch.Tensor],

position_ids: Optional[torch.Tensor] = None,

attention_mask: Optional[torch.Tensor] = None,

labels: Optional[torch.Tensor] = None,

loss_mask: Optional[torch.Tensor] = None,

padding_mask: Optional[torch.Tensor] = None,

inference_context=None,

runtime_gather_output: Optional[bool] = None,

packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,

images: Optional[torch.Tensor] = None,

pixel_values: Optional[torch.Tensor] = None,

imgs_sizes: Optional[torch.Tensor] = None,

vision_packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None,

num_frames: Optional[torch.Tensor] = None,

sound_clips: Optional[torch.Tensor] = None,

sound_length: Optional[torch.Tensor] = None,

*,

media_token_validity_mask: torch.Tensor | None = None,

inference_params=None,

**kwargs,

) → torch.Tensor | tuple[torch.Tensor, torch.Tensor]

Insert media into the expanded sequence, shard for CP, then call NemotronH.

Returns:

Model output, or (output, local_loss_mask) when this model applies a context-parallel shard to the supervision tensors.