[inference] fix: Support MCore dev inference mode

src/megatron/bridge/inference/vlm/_mcore_compat.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compatibility helpers for Megatron-Core inference API differences."""
+
+try:
+    from megatron.core.inference.utils import InferenceMode
+except ImportError as exc:
+    if "InferenceMode" not in str(exc):
+        raise
+
+    # TODO(mcore-dev): remove this guard when Megatron-Core dev exposes InferenceMode from
+    # megatron.core.inference.utils.
+    class InferenceMode:
+        """No-op compatibility shim for MCore commits without InferenceMode."""
+
+        @classmethod
+        def is_active(cls) -> bool:
+            """Return whether MCore's process-wide inference mode is active."""
+            return False
+
+        @classmethod
+        def set_active(cls) -> None:
+            """Mark inference as active when the backing MCore API exists."""
+
+        @classmethod
+        def unset_active(cls) -> None:
+            """Mark inference as inactive when the backing MCore API exists."""

src/megatron/bridge/inference/vlm/vlm_engine.py

@@ -18,9 +18,10 @@
 from megatron.core.inference.engines.static_engine import StaticInferenceEngine
 from megatron.core.inference.inference_request import InferenceRequest
 from megatron.core.inference.sampling_params import SamplingParams
-from megatron.core.inference.utils import InferenceMode
 from PIL.Image import Image
 
+from ._mcore_compat import InferenceMode
+
 
 class VLMEngine(StaticInferenceEngine):
     """VLM inference engine extending MCoreEngine with image support."""

src/megatron/bridge/models/gpt_provider.py

@@ -268,11 +268,6 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGP
         if self.init_model_with_meta_device:
             model_init_device_context = partial(torch.device, device="meta")
 
-        # Guard for main/dev branch submodule compat: mtp_block_spec was added in the dev branch.
-        # TODO: remove guard once the addition lands in main and Bridge pins the new main commit.
-        kwargs = {}
-        if "mtp_block_spec" in inspect.signature(MCoreGPTModel.__init__).parameters:
-            kwargs["mtp_block_spec"] = mtp_block_spec(self, vp_stage=vp_stage)
         if self.attention_backend == AttnBackend.local:
             if hasattr(transformer_layer_spec, "submodules"):
                 transformer_layer_spec.submodules.self_attention.submodules.core_attention = MCoreDotProductAttention
@@ -308,7 +303,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGP
                 scatter_embedding_sequence_parallel=self.scatter_embedding_sequence_parallel,
                 pg_collection=self._pg_collection,
                 vp_stage=vp_stage,
-                **kwargs,
+                mtp_block_spec=mtp_block_spec(self, vp_stage=vp_stage),
             )
 
         # If using full TE layer, need to set TP, CP group since the module call

src/megatron/bridge/models/mamba/mamba_provider.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import inspect
 import logging
 import warnings
 from dataclasses import dataclass, field
@@ -29,7 +28,7 @@
 from megatron.core.pipeline_parallel.utils import is_pp_first_stage, is_pp_last_stage
 from megatron.core.post_training.modelopt.mamba.model_specs import get_mamba_stack_modelopt_spec
 from megatron.core.process_groups_config import ProcessGroupCollection
-from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, parse_hybrid_pattern
+from megatron.core.ssm.mamba_hybrid_layer_allocation import Symbols, get_hybrid_total_layer_count, parse_hybrid_pattern
 from megatron.core.transformer import ModuleSpec
 from megatron.core.transformer.enums import AttnBackend
 
@@ -40,51 +39,8 @@
 from megatron.bridge.utils.vocab_utils import calculate_padded_vocab_size
 
 
-try:
-    from megatron.core.ssm.mamba_hybrid_layer_allocation import (
-        get_hybrid_total_layer_count as _mcore_get_hybrid_total_layer_count,
-    )
-except ImportError:
-    # TODO(yuya): remove fallback once MCore pin includes get_hybrid_total_layer_count
-    _mcore_get_hybrid_total_layer_count = None
-
-# MCore renamed `hybrid_override_pattern` → `hybrid_layer_pattern` in the dev branch.
-# Support both main and dev branch submodule by detecting which parameter is present at import time.
-# TODO: remove fallback once the dev rename lands in main and Bridge pins the new main commit.
-_MCORE_MAMBA_INIT_PARAMS = set(inspect.signature(MCoreMambaModel.__init__).parameters)
-_HYBRID_LAYER_PATTERN_KWARG = (
-    "hybrid_layer_pattern" if "hybrid_layer_pattern" in _MCORE_MAMBA_INIT_PARAMS else "hybrid_override_pattern"
-)
-
-
 logger = logging.getLogger(__name__)
 
-_HYBRID_MAIN_PATTERN_SYMBOLS = frozenset({"M", "*", "-", "E", "|"})
-
-
-def _fallback_get_hybrid_total_layer_count(pattern: str) -> int:
-    """Count main-decoder layers for older MCore branches.
-
-    Older MCore revisions predate ``get_hybrid_total_layer_count`` and do not
-    understand pipe-delimited fVPP layouts. Bridge still needs to derive
-    ``num_layers`` correctly for both legacy and newer hybrid patterns.
-    """
-
-    main_pattern = pattern.split("/")[0]
-    invalid_chars = sorted({char for char in main_pattern if char not in _HYBRID_MAIN_PATTERN_SYMBOLS})
-    if invalid_chars:
-        raise ValueError(
-            f"In main pattern, '{invalid_chars[0]}' is not a valid layer symbol. "
-            f"Valid symbols are: {_HYBRID_MAIN_PATTERN_SYMBOLS}"
-        )
-    return len(main_pattern.replace("|", ""))
-
-
-def _get_hybrid_total_layer_count(pattern: str) -> int:
-    if _mcore_get_hybrid_total_layer_count is not None:
-        return _mcore_get_hybrid_total_layer_count(pattern)
-    return _fallback_get_hybrid_total_layer_count(pattern)
-
 
 def modelopt_mamba_stack_spec(config: "MambaModelProvider") -> ModuleSpec:
     """Mamba stack specification for quantization with ModelOpt.
@@ -257,7 +213,7 @@ def finalize(self) -> None:
         # Check if hybrid_layer_pattern is specified and derive num_layers from pattern
         if self.hybrid_layer_pattern is not None:
             # Derive num_layers from pattern
-            num_layers_in_pattern = _get_hybrid_total_layer_count(self.hybrid_layer_pattern)
+            num_layers_in_pattern = get_hybrid_total_layer_count(self.hybrid_layer_pattern)
             if self.num_layers is not None:
                 if used_hybrid_override_pattern:
                     assert self.num_layers == num_layers_in_pattern, (
@@ -314,7 +270,7 @@ def provide(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreMa
             mamba_stack_spec=mamba_stack_spec,
             vocab_size=padded_vocab_size,
             max_sequence_length=self.seq_length,
-            **{_HYBRID_LAYER_PATTERN_KWARG: self.hybrid_layer_pattern},
+            hybrid_layer_pattern=self.hybrid_layer_pattern,
             fp16_lm_cross_entropy=self.fp16_lm_cross_entropy,
             parallel_output=self.parallel_output,
             share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,

src/megatron/bridge/models/qwen3_asr/hf_qwen3_asr/modeling_qwen3_asr.py

@@ -993,6 +993,10 @@ def forward(
         cache_position: Optional[torch.LongTensor] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Union[tuple, BaseModelOutputWithPast]:
+        r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence.
+        """
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
@@ -1172,6 +1176,8 @@ def forward(
         **kwargs,
     ) -> Union[tuple, Qwen3ASRThinkerCausalLMOutputWithPast]:
         r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the input sequence tokens in the sequence.
         feature_attention_mask (`torch.Tensor` of shape `(batch_size, feature_sequence_length)`, *optional*):
             Mask to avoid performing attention on padding feature indices. Mask values selected in `[0, 1]`:
             - 1 for tokens that are **not masked**,

src/megatron/bridge/recipes/utils/optimizer_utils.py

@@ -12,19 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import dataclasses
 from typing import Optional
 
 from megatron.bridge.training.config import OptimizerConfig, SchedulerConfig
 
 
-# MCore renamed `muon_use_nesterov` → `muon_nesterov` in the dev branch.
-# Support both main and dev branch submodule by detecting which field is present at import time.
-# TODO: remove fallback once the dev rename lands in main and Bridge pins the new main commit.
-_OPTIMIZER_CONFIG_FIELDS = {f.name for f in dataclasses.fields(OptimizerConfig)}
-_MUON_NESTEROV_KWARG = "muon_nesterov" if "muon_nesterov" in _OPTIMIZER_CONFIG_FIELDS else "muon_use_nesterov"
-
-
 def distributed_muon_with_cosine_annealing(
     precision: str = "bf16-mixed",
     muon_momentum: float = 0.95,
@@ -91,7 +83,7 @@ def distributed_muon_with_cosine_annealing(
         bf16=precision == "bf16-mixed",
         fp16=precision == "16-mixed",
         muon_momentum=muon_momentum,
-        **{_MUON_NESTEROV_KWARG: muon_use_nesterov},
+        muon_nesterov=muon_use_nesterov,
         muon_scale_mode=muon_scale_mode,
         muon_fp32_matmul_prec=muon_fp32_matmul_prec,
         muon_num_ns_steps=muon_num_ns_steps,

src/megatron/bridge/training/initialize.py

@@ -789,8 +789,8 @@ def _initialize_distributed(
         if parallel_state.model_parallel_is_initialized():
             print("model parallel is already initialized")
         else:
-            # Guard for main/dev branch submodule compat: hybrid_context_parallel was added in the dev branch.
-            # TODO: remove guard once the addition lands in main and Bridge pins the new main commit.
+            # Guard for main/dev branch submodule compat: dev exposes dynamic_context_parallel instead.
+            # TODO(mcore-dev): remove once dev exposes hybrid_context_parallel or Bridge migrates the config.
             _init_mp_params = set(inspect.signature(parallel_state.initialize_model_parallel).parameters)
             _optional_kwargs = {}
             if "hybrid_context_parallel" in _init_mp_params:

src/megatron/bridge/training/optim.py

@@ -19,20 +19,8 @@
     MegatronOptimizer,
     OptimizerConfig,
     get_megatron_optimizer,
+    get_mup_config_overrides,
 )
-
-
-# TODO: Remove try/except once `get_mup_config_overrides` lands in mcore main.
-#       This guard exists because the symbol lives in mcore dev but not yet in
-#       the main branch that the submodule tracks.
-#
-#       We assign None (not a bool flag) so the module attribute always exists
-#       and tests can patch it without AttributeError.
-try:
-    from megatron.core.optimizer import get_mup_config_overrides
-except ImportError:
-    get_mup_config_overrides = None  # type: ignore[assignment]
-
 from megatron.core.optimizer.muon import get_megatron_muon_optimizer
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.process_groups_config import ProcessGroupCollection
@@ -78,11 +66,9 @@ def setup_optimizer(
     )
 
     # Apply μP optimizer scaling if enabled on the model config.
-    # Guard on the callable itself (None when mcore main lacks the symbol) so
-    # unit tests can patch the module attribute without hitting AttributeError.
     model_chunks = model if isinstance(model, list) else [model]
     model_config = get_model_config(model_chunks[0])
-    if get_mup_config_overrides is not None and getattr(model_config, "use_mup", False):
+    if getattr(model_config, "use_mup", False):
         mup_overrides = get_mup_config_overrides(
             config=optimizer_config,
             mup_width_mult=model_config.mup_width_mult,

src/megatron/bridge/training/state.py

@@ -20,25 +20,14 @@
 from typing import Any, Optional
 
 import torch
+from megatron.core.dist_checkpointing.strategies.torch import get_async_strategy
+from megatron.core.energy_monitor import EnergyMonitor
+from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.timers import Timers
 from megatron.core.utils import StragglerDetector
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.utils.tensorboard.writer import SummaryWriter
 
-
-# TODO: Remove try/except guards once these land in mcore dev.
-try:
-    from megatron.core.dist_checkpointing.strategies.torch import get_async_strategy
-except ImportError:
-    get_async_strategy = None  # type: ignore[assignment]
-
-try:
-    from megatron.core.energy_monitor import EnergyMonitor
-except ImportError:
-    EnergyMonitor = None  # type: ignore[assignment]
-
-from megatron.core.process_groups_config import ProcessGroupCollection
-
 from megatron.bridge.training.config import ConfigContainer
 from megatron.bridge.training.nvrx_straggler import NVRxStragglerDetectionManager
 from megatron.bridge.training.tokenizers.tokenizer import build_tokenizer
@@ -412,21 +401,9 @@ def initialize_async_checkpoint_worker(self) -> None:
             and self.cfg.checkpoint.save is not None
             and self.cfg.checkpoint.async_save
         ):
-            if get_async_strategy is not None:
-                # mcore main path: get_async_strategy selects nvrx vs mcore backend
-                async_strategy, async_modules = get_async_strategy(self.cfg.checkpoint.async_strategy)
-                async_calls_queue_cls = async_modules["AsyncCallsQueue"]
-                get_write_results_queue_fn = async_modules["get_write_results_queue"]
-            else:
-                # mcore dev path: nvrx modules merged into core, no strategy selector
-                from megatron.core.dist_checkpointing.strategies.async_utils import AsyncCallsQueue
-                from megatron.core.dist_checkpointing.strategies.filesystem_async import (
-                    get_write_results_queue,
-                )
-
-                async_strategy = None
-                async_calls_queue_cls = AsyncCallsQueue
-                get_write_results_queue_fn = get_write_results_queue
+            async_strategy, async_modules = get_async_strategy(self.cfg.checkpoint.async_strategy)
+            async_calls_queue_cls = async_modules["AsyncCallsQueue"]
+            get_write_results_queue_fn = async_modules["get_write_results_queue"]
 
             self._async_calls_queue = async_calls_queue_cls(persistent=self.cfg.checkpoint.use_persistent_ckpt_worker)
 
@@ -466,7 +443,6 @@ def energy_monitor(self) -> Optional[Any]:
             and self._energy_monitor is None
             and self.cfg is not None
             and self.cfg.logger.log_energy
-            and EnergyMonitor is not None
         ):
             self._energy_monitor = EnergyMonitor()
             self._energy_monitor_created = True

src/megatron/bridge/training/tokenizers/tokenizer.py

@@ -78,38 +78,14 @@ def build_tokenizer(config: TokenizerConfig, **kwargs) -> MegatronTokenizer:
         tokenizer_library = "null-text"
         if config.vocab_size:
             kwargs["vocab_size"] = config.vocab_size
-        # TODO(mcore-guard): Remove try/except once mcore main and dev both support
-        # "null-text"/"null-multimodal" tokenizer library names (dev renamed "null" → split names).
-        try:
-            metadata = {"library": tokenizer_library}
-            tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs)
-        except AssertionError:
-            # Legacy mcore exposed NullTokenizer under library "null" and internally reserved the
-            # top id for the pad token, requiring callers to pass vocab_size - 1 to obtain the
-            # requested effective vocab size.
-            metadata = {"library": "null"}
-            if "vocab_size" in kwargs:
-                kwargs["vocab_size"] = kwargs["vocab_size"] - 1
-            tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs)
-
-        return tokenizer
+        return MegatronTokenizer.from_pretrained(metadata_path={"library": tokenizer_library}, **kwargs)
     elif config.tokenizer_type == "NullMultimodalTokenizer":
         # NullMultimodalTokenizer still reserves the top id for the pad token, so the effective
-        # vocab size is passed as vocab_size - 1 under both the new "null-multimodal" and the
-        # legacy "null" library names.
+        # vocab size is passed as vocab_size - 1.
         tokenizer_library = "null-multimodal"
         if config.vocab_size:
             kwargs["vocab_size"] = config.vocab_size - 1
-        # TODO(mcore-guard): Remove try/except once mcore main and dev both support
-        # "null-text"/"null-multimodal" tokenizer library names (dev renamed "null" → split names).
-        try:
-            metadata = {"library": tokenizer_library}
-            tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs)
-        except AssertionError:
-            metadata = {"library": "null"}
-            tokenizer = MegatronTokenizer.from_pretrained(metadata_path=metadata, **kwargs)
-
-        return tokenizer
+        return MegatronTokenizer.from_pretrained(metadata_path={"library": tokenizer_library}, **kwargs)
 
     if config.metadata_path:
         metadata = config.metadata_path

tests/unit_tests/inference/vlm/test_vlm_engine.py

@@ -15,8 +15,8 @@
 from unittest.mock import MagicMock
 
 from megatron.core.inference.contexts import StaticInferenceContext
-from megatron.core.inference.utils import InferenceMode
 
+from megatron.bridge.inference.vlm._mcore_compat import InferenceMode
 from megatron.bridge.inference.vlm.vlm_engine import VLMEngine