bridge.models.nemotron_omni.nemotron_omni_bridge#

Nemotron Omni conversion bridges.

Standalone bridge for the Nemotron-3 Omni family (HF architecture NemotronH_Nano_Omni_Reasoning_V3). Inherits the language / vision / mamba parameter mappings from :class:NemotronVLBridge and adds:

  • Omni-specific CONFIG_MAPPING entries (Mamba shape fields used by the hybrid LLM and the MoE shared-expert intermediate size).

  • An overridden :meth:provider_bridge that produces a

    class:

    NemotronOmniModelProvider (MoE language model + RADIO ViT vision

    • optional Parakeet sound encoder) instead of the dense VL provider.

  • A :meth:mapping_registry override that adds the temporal video_embedder parameter and the sound projection / sound encoder parameters (the latter via a single ** wildcard, since the Megatron sound encoder is HF transformers’ ParakeetEncoder and the parameter names line up 1:1 with sound_encoder.encoder.*).

  • ADDITIONAL_FILE_PATTERNS covering the bespoke Omni HF modeling / processing / audio files that need to be copied during HF export.

Module Contents#

Classes#

_LegacyNemotronOmniBridge

Load MoE Nano Omni checkpoints carrying the historical V2 label.

NemotronOmniBridge

Bridge for the canonical expanded-sequence Nemotron-3 Omni model.

Nemotron35SuperVLBridge

Bridge for Nemotron 3.5 Super VL using the shared Omni media stack.

NemotronOmniLlavaBridge

Deprecated fallback bridge for the historical collapse/expand model.

Functions#

_copy_mapping_with_prefixes

Copy a mapping while preserving its conversion implementation.

nemotron_omni_mimo_conversion_spec

Reuse Omni weight mappings with separate language and image grids.

API#

bridge.models.nemotron_omni.nemotron_omni_bridge._copy_mapping_with_prefixes(
mapping,
*,
megatron_prefix: str,
hf_prefix: str,
)#

Copy a mapping while preserving its conversion implementation.

class bridge.models.nemotron_omni.nemotron_omni_bridge._LegacyNemotronOmniBridge#

Bases: megatron.bridge.models.nemotron_vl.nemotron_vl_bridge.NemotronVLBridge

Load MoE Nano Omni checkpoints carrying the historical V2 label.

class bridge.models.nemotron_omni.nemotron_omni_bridge.NemotronOmniBridge#

Bases: megatron.bridge.models.nemotron_vl.nemotron_vl_bridge.NemotronVLBridge

Bridge for the canonical expanded-sequence Nemotron-3 Omni model.

_HF_DYNAMIC_MODULE_IMPORTS = <Multiline-String>#
_HF_PASSTHROUGH_KEYS#

(‘sound_encoder.encoder.feature_extractor.featurizer.fb’, ‘sound_encoder.encoder.feature_extractor.f…

CONFIG_MAPPING#

None

ADDITIONAL_FILE_PATTERNS#

[‘modeling*.py’, ‘configuration*.py’, ‘processing*.py’, ‘processing_utils.py’, ‘image_processing*.py…

postprocess_hf_export_artifacts(path: pathlib.Path) → None#

Make transitive Omni configuration modules discoverable on local reload.

The pinned Nemotron Omni Hugging Face repositories expose modeling.py as their auto_map entrypoint. Fail explicitly if that required export artifact is absent so an incomplete checkpoint is never reported as saved.

provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.models.nemotron_omni.nemotron_omni_provider.NemotronOmniModelProvider#

Create a NemotronOmniModelProvider from the HF Omni config.

Always returns an Omni provider (MoE language model + RADIO ViT vision + optional Parakeet sound encoder). The presence of sound_config is the Hugging Face checkpoint’s sound capability.

classmethod megatron_to_hf_config(provider) → dict#

Export sound capability consistently with model construction.

_mtp_hf_prefix() → str#

Return the HF prefix applied to Nemotron-H MTP weights.

_llava_mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#

Build mappings for the historical LLaVA wrapper namespace.

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#

Return top-level media mappings plus prefixed NemotronH mappings.

stream_weights_megatron_to_hf(
megatron_model: megatron.bridge.models.nemotron_omni.modeling_nemotron_omni.NemotronOmniModel | list[megatron.bridge.models.nemotron_omni.modeling_nemotron_omni.NemotronOmniModel],
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
cpu: bool = True,
show_progress: bool = True,
conversion_tasks: list[megatron.bridge.models.conversion.model_bridge.WeightConversionTask] | None = None,
merge_adapter_weights: bool = True,
weight_dtype: torch.dtype | None = None,
with_megatron_names: bool = False,
current_pp_stage_only: bool = False,
) → collections.abc.Iterable[megatron.bridge.models.conversion.model_bridge.HFWeightTuple | megatron.bridge.models.conversion.model_bridge.HFSourcedWeightTuple]#

Export model weights and preserve immutable source-only buffers.

class bridge.models.nemotron_omni.nemotron_omni_bridge.Nemotron35SuperVLBridge#

Bases: bridge.models.nemotron_omni.nemotron_omni_bridge.NemotronOmniBridge

Bridge for Nemotron 3.5 Super VL using the shared Omni media stack.

_HF_SUMMARY_IDXS_BUFFER#

‘vision_model.summary_idxs’

_HF_PASSTHROUGH_KEYS#

()

_HF_SHARED_MTP_BLOCKS#

1

_MCORE_MTP_PREDICTION_DEPTHS#

2

classmethod _validate_shared_mtp_config(llm_config) → None#

Validate the serialized shared block before choosing training depths.

text_only_pretrained(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM#

Select the native Nemotron-H language checkpoint, excluding all media.

Preserve the HF count of one serialized shared MTP block. Super text recipes explicitly set two training prediction depths, independently of this checkpoint representation.

postprocess_hf_export_artifacts(path: pathlib.Path) → None#

Require the direct Transformers entrypoint used by Super VL exports.

provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.models.nemotron_omni.nemotron_omni_provider.NemotronOmniModelProvider#

Create the shared Omni provider with Super-VL checkpoint features enabled.

classmethod megatron_to_hf_config(provider) → dict#

Preserve HF’s single-block serialization for the repeated MTP head.

postprocess_hf_export_weights(path: pathlib.Path) → None#

Add the deterministic RADIO summary buffer omitted by the source index.

_mtp_hf_prefix() → str#

Nemotron 3.5 Super VL nests MTP below language_model.

mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#

Add the Super-VL vision final norm to the shared Omni mappings.

class bridge.models.nemotron_omni.nemotron_omni_bridge.NemotronOmniLlavaBridge#

Bases: bridge.models.nemotron_omni.nemotron_omni_bridge.NemotronOmniBridge

Deprecated fallback bridge for the historical collapse/expand model.

Use :class:NemotronOmniBridge, which is the canonical AutoBridge registration and consumes processor-expanded media-token sequences.

provider_bridge(
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
) → megatron.bridge.models.nemotron_omni.nemotron_omni_provider.NemotronOmniLlavaModelProvider#
mapping_registry() → megatron.bridge.models.conversion.mapping_registry.MegatronMappingRegistry#
bridge.models.nemotron_omni.nemotron_omni_bridge.nemotron_omni_mimo_conversion_spec(
source_bridge: bridge.models.nemotron_omni.nemotron_omni_bridge.NemotronOmniBridge,
hf_pretrained: megatron.bridge.models.hf_pretrained.causal_lm.PreTrainedCausalLM,
parallelism_config: megatron.bridge.models.megatron_mimo.megatron_mimo_config.MegatronMIMOParallelismConfig,
) → tuple[megatron.bridge.models.megatron_mimo.megatron_mimo_provider.MegatronMIMOProvider, list[megatron.bridge.models.megatron_mimo.conversion.MIMOComponent]]#

Reuse Omni weight mappings with separate language and image grids.

The RADIO encoder and projector have separate routes sharing the image component’s process groups.