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

`NemotronOmniMimoRadioEncoder`	RADIO encoder returning one pixel-shuffled row per Nemotron image token.
`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.
`_pad_patch_grid_to_even`	Zero-pad a patch grid to even extents so pixel shuffle can consume it.
`_postprocess_radio_embeddings`	Remove per-image class tokens and pixel-shuffle dynamic RADIO patch grids.
`_project_multimodal_embeddings`	Project media rows, padding only the temporary FP8 compute input.
`patchify_dynamic_images`	Convert padded processor pixels to RADIO’s packed patch representation.

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.

bridge.models.nemotron_omni.modeling_nemotron_omni._pad_patch_grid_to_even( features: torch.Tensor, *, height: int, width: int, ) → tuple[torch.Tensor, int, int]#

Zero-pad a patch grid to even extents so pixel shuffle can consume it.

Only the 1x1 placeholder images that the context-parallel split injects to keep every rank non-empty reach this path; real Omni frames are already a multiple of two in both dimensions.

bridge.models.nemotron_omni.modeling_nemotron_omni._postprocess_radio_embeddings( encoded: torch.Tensor, imgs_sizes: torch.Tensor, *, patch_dim: int, class_tokens: int, pad_patch_grid: bool = False, ) → torch.Tensor#

Remove per-image class tokens and pixel-shuffle dynamic RADIO patch grids.

Context-parallel shards may contain odd-sized placeholder grids, which require padding before pixel shuffle. Unsharded inputs retain strict grid validation.

bridge.models.nemotron_omni.modeling_nemotron_omni._project_multimodal_embeddings( projection: torch.nn.Module, embeddings: torch.Tensor, ) → torch.Tensor#: Project media rows, padding only the temporary FP8 compute input.

bridge.models.nemotron_omni.modeling_nemotron_omni.patchify_dynamic_images( images: torch.Tensor, imgs_sizes: torch.Tensor, *, patch_dim: int, ) → torch.Tensor#

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

Accepts [num_images, channels, padded_height, padded_width] pixels and returns [1, total_patches, channels * patch_dim**2]. Already-patchified inputs ([total_patches, C*P*P] or [1, total_patches, C*P*P]) pass through unchanged.

class bridge.models.nemotron_omni.modeling_nemotron_omni.NemotronOmniMimoRadioEncoder#

Bases: megatron.core.models.vision.radio.RADIOViTModel

RADIO encoder returning one pixel-shuffled row per Nemotron image token.

Output is [num_image_tokens, 4 * hidden_size]: the matching MultimodalProjector input projection maps it to the language hidden size.

forward( pixel_values: Optional[torch.Tensor] = None, imgs_sizes: Optional[torch.Tensor] = None, num_frames: Optional[torch.Tensor] = None, vision_packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, *, images: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, ) → torch.Tensor#

Encode dynamic-resolution images (or temporal tubelets) into projected-ready rows.

Parameters:

pixel_values – Padded pixels [N, C, H, W] or packed patches [1, total_patches, C*P*P].
imgs_sizes – Per-image (height, width) after resize, [N, 2].
num_frames – Frames per media item; required when the encoder has a temporal patch dimension greater than one.
vision_packed_seq_params – Optional THD boundaries; derived from imgs_sizes when omitted.
images – Alias for pixel_values.
attention_mask – Unused; accepted for encoder-interface compatibility.

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, vision_dp_over_cp: 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

_expose_language_model_for_cuda_graph_helper() → None#

Expose language-model fields on the multimodal root when CUDA graphs are enabled.

MCore’s CUDA graph helper discovers decoder on the top-level model, while Nemotron Omni stores it below language_model. Properties expose the nested modules without registering duplicate root aliases or changing checkpoint keys.

property rotary_pos_emb#: Expose the nested language model’s rotary embeddings to MCore helpers.

property decoder#: Expose the nested language decoder without registering a module alias.

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

_split_images_across_context_parallel_ranks( images: torch.Tensor, imgs_sizes: torch.Tensor, vision_packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams, *, num_frames: Optional[torch.Tensor], temporal_patch_size: int, )#: Give each CP rank an integer share of the patched images or tubelets.

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

bridge.models.nemotron_omni.modeling_nemotron_omni#

Module Contents#

Classes#

Functions#

API#

`bridge.models.nemotron_omni.modeling_nemotron_omni`#