`bridge.models.nemotron_omni.nemotron_omni_provider`#

Module Contents#

Classes#

`NemotronVLModelProvider`	Abstract base provider for Nemotron VL model variants.
`_NemotronOmniModelProviderBase`	Shared configuration and builders for Nemotron Omni models.
`NemotronOmniModelProvider`	Provider for the canonical expanded-sequence Nemotron Omni model.
`NemotronOmniLlavaModelProvider`	Deprecated fallback provider for the historical collapse/expand path.

Functions#

_get_transformer_engine_projection_submodules

Build the standard dense TE MLP submodules used by media projectors.

Data#

`NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT`
`NEMOTRON_OMNI_LLAVA_CONTRACT`
`_IMAGES_MODALITY_KEY`
`_RADIO_ENCODER_KEY`

API#

bridge.models.nemotron_omni.nemotron_omni_provider.NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT#: ‘expanded_sequence_v1’

bridge.models.nemotron_omni.nemotron_omni_provider.NEMOTRON_OMNI_LLAVA_CONTRACT#: ‘llava_collapse_expand_v1’

bridge.models.nemotron_omni.nemotron_omni_provider._IMAGES_MODALITY_KEY#: ‘images’

bridge.models.nemotron_omni.nemotron_omni_provider._RADIO_ENCODER_KEY#: ‘radio’

bridge.models.nemotron_omni.nemotron_omni_provider._get_transformer_engine_projection_submodules()#: Build the standard dense TE MLP submodules used by media projectors.

class bridge.models.nemotron_omni.nemotron_omni_provider.NemotronVLModelProvider#

Bases: megatron.bridge.models.hybrid.hybrid_provider.HybridModelProvider, abc.ABC

Abstract base provider for Nemotron VL model variants.

Provides common VL fields, RADIO ViT-H vision config building methods, and vision projection config building methods shared by dense and MoE variants. Concrete subclasses set LLM-specific defaults (hidden_size, hybrid pattern, etc.) and may override provide() for variant-specific assembly.

mamba_num_groups: int#: 8

num_query_groups: int#: 8

make_vocab_size_divisible_by: int#: 128

activation_func: Callable#: None

masked_softmax_fusion: bool#: True

apply_query_key_layer_scaling: bool#: False

persist_layer_norm: bool#: True

first_last_layers_bf16: bool#: True

is_hybrid_model: bool#: True

moe_aux_loss_coeff: float#: 0.0001

moe_router_score_function: str#: ‘sigmoid’

moe_router_enable_expert_bias: bool#: True

moe_router_load_balancing_type: str#: ‘seq_aux_loss’

moe_router_dtype: str#: ‘fp32’

moe_grouped_gemm: bool#: True

moe_token_dispatcher_type: str#: ‘alltoall’

moe_permute_fusion: bool#: True

moe_shared_expert_overlap: bool#: True

scatter_embedding_sequence_parallel: bool#: False

attention_softmax_in_fp32: bool#: True

vision_model_type: str#: ‘radio’

language_model_type: str = <Multiline-String>#

image_token_index: int#: 0

img_start_token_id: int#: 0

img_end_token_id: int#: 0

tokenizer_type: str = <Multiline-String>#

dynamic_resolution: bool#: False

use_vision_backbone_fp8_arch: bool#: True

radio_force_eval_mode: bool#: True

radio_force_cpe_eval_mode: bool#: True

radio_interpolate_only_cpe: bool#: True

radio_cpe_aspect_ratio_select: bool#: False

radio_disable_cpe: bool#: False

recompute_vision: bool#: False

vision_recompute_granularity: Literal[full, selective] | None#: None

vision_recompute_modules: list[str] | None#: None

vision_recompute_method: Literal[uniform, block] | None#: None

vision_recompute_num_layers: int | None#: None

vision_proj_ffn_hidden_size: int#: 20480

vision_class_token_len: Optional[int]#: None

freeze_language_model: bool#: False

freeze_vision_model: bool#: False

freeze_vision_projection: bool#: False

_build_vision_config(language_cfg)#: Build RADIO ViT-H vision encoder config from a language config copy.

_build_vision_projection_config(language_cfg)#: Build vision projection MLP config from a language config copy.

provide_language_model( pre_process=None, post_process=None, vp_stage=None, )#

class bridge.models.nemotron_omni.nemotron_omni_provider._NemotronOmniModelProviderBase#

Bases: bridge.models.nemotron_omni.nemotron_omni_provider.NemotronVLModelProvider

Shared configuration and builders for Nemotron Omni models.

Extends NemotronVLModelProvider with sound-specific fields and fixes RADIO to the public model’s dynamic-resolution input contract.

dynamic_resolution: Literal[True]#: True

has_sound: bool#: False

sound_model_type: str#: ‘parakeet’

sound_hidden_size: int#: 1024

sound_projection_hidden_size: int#: 4096

sound_context_token_id: int#: 0

sound_config: Optional[dict]#: None

freeze_sound_encoder: bool#: False

freeze_sound_projection: bool#: False

temporal_patch_dim: int#: 1

separate_video_embedder: bool#: False

temporal_ckpt_compat: bool#: False

vision_final_layernorm: bool#: False

vision_dp_over_cp: bool#: True

nemotron_omni_contract: Optional[str]#: None

_validate_omni_config() → None#

finalize() → None#: Finalize a dynamic-resolution Nemotron Omni provider.

_require_model_contract(expected: str, implementation: str) → None#

_build_vision_config(language_cfg)#

Pin vision encoder to PP=1 (Omni training uses PP>1 on the LLM).

The dense Nemotron-VL recipe runs with PP=1 everywhere, so the base VL provider doesn’t pin this; for Omni we always co-locate the vision encoder with the first PP stage.

_build_vision_projection_config(language_cfg)#

Build the vision projection MLP config.

Nemotron Omni applies squared ReLU in the vision projection MLP, matching the HF checkpoint and vLLM implementation. Also pin the projector to PP=1 (see :meth:_build_vision_config).

_build_sound_projection_config(language_cfg)#: Build sound projection config (mirrors _build_vision_projection_config).

_build_sound_encoder()#: Build BridgeSoundEncoder from sound_config dict.

_build_sound_modules(language_cfg, *, add_encoder: bool)#: Build optional sound modules on the encoder pipeline stage.

_provide_llava(pre_process=None, post_process=None, vp_stage=None)#

Assemble the legacy LLaVA collapse/expand implementation.

Duplicates the VL provide() logic because LLaVAModel requires sound kwargs at construction time – they can’t be added after. This is intentional to maintain zero changes to nemotron_vl/.

_configure_llava_preprocess_contract( llava_model: megatron.core.models.multimodal.llava_model.LLaVAModel, ) → None#: Align MCore preprocessing with the legacy collator’s replacement counts.

class bridge.models.nemotron_omni.nemotron_omni_provider.NemotronOmniModelProvider#

Bases: bridge.models.nemotron_omni.nemotron_omni_provider._NemotronOmniModelProviderBase

Provider for the canonical expanded-sequence Nemotron Omni model.

mimo_supports_mtp: ClassVar[bool]#: True

radio_interpolate_only_cpe: bool#: False

property special_token_ids: dict[str, int]#: Return the MIMO modality placeholder token ids.

build_language_model_spec( pp_rank: Optional[int] = 0, ) → megatron.core.transformer.spec_utils.ModuleSpec#

Build the MIMO language module spec (Nemotron-H hybrid decoder with MTP).

Mirrors :meth:HybridModelProvider.provide; pre_process, post_process and pg_collection are injected by the MIMO provider.

_build_mimo_media_config(builder: Callable)#

Derive an encoder/projector config copy that survives MIMO re-finalization.

The standard model derives these copies from the finalized language config and never finalizes them again. The MIMO provider finalizes every config it finds in a spec, so the hybrid/MTP pattern fields that would re-derive the language layer count are cleared and expert parallelism is made dense.

build_vision_encoder_spec() → megatron.core.transformer.spec_utils.ModuleSpec#: Build the RADIO encoder spec for the MIMO images modality.

build_vision_input_projection_spec() → megatron.core.transformer.spec_utils.ModuleSpec#: Build the pixel-shuffled RADIO feature to language hidden projector spec.

build_mimo_modality_submodules_spec() → dict[str, megatron.core.transformer.spec_utils.ModuleSpec]#: Build the MIMO modality specs: RADIO encoder plus vision projector.

validate_model_contract() → None#: Reject ambiguous or legacy serialized provider configurations.

provide(pre_process=None, post_process=None, vp_stage=None)#: Build the top-level model from the shared Nano/Super fields.

class bridge.models.nemotron_omni.nemotron_omni_provider.NemotronOmniLlavaModelProvider#

Bases: bridge.models.nemotron_omni.nemotron_omni_provider.NemotronOmniModelProvider

Deprecated fallback provider for the historical collapse/expand path.

Use :class:NemotronOmniModelProvider, which constructs the canonical processor-expanded model with collator-owned packing.

radio_interpolate_only_cpe: bool#: True

validate_model_contract() → None#: Require an explicitly selected legacy LLaVA contract.

provide(pre_process=None, post_process=None, vp_stage=None)#

bridge.models.nemotron_omni.nemotron_omni_provider#