From 3336f64927feff1f33a1e3f179694651884e9c78 Mon Sep 17 00:00:00 2001 From: Yu Yao Date: Mon, 18 May 2026 11:43:00 -0700 Subject: [PATCH 1/2] [inference] fix: Support MCore dev inference mode Signed-off-by: Yu Yao --- .../bridge/inference/vlm/_mcore_compat.py | 38 +++++++++++++++++++ .../bridge/inference/vlm/vlm_engine.py | 3 +- .../hf_qwen3_asr/modeling_qwen3_asr.py | 6 +++ .../inference/vlm/test_vlm_engine.py | 2 +- 4 files changed, 47 insertions(+), 2 deletions(-) create mode 100644 src/megatron/bridge/inference/vlm/_mcore_compat.py diff --git a/src/megatron/bridge/inference/vlm/_mcore_compat.py b/src/megatron/bridge/inference/vlm/_mcore_compat.py new file mode 100644 index 0000000000..eccc00dc61 --- /dev/null +++ b/src/megatron/bridge/inference/vlm/_mcore_compat.py @@ -0,0 +1,38 @@ +# 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: 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.""" diff --git a/src/megatron/bridge/inference/vlm/vlm_engine.py b/src/megatron/bridge/inference/vlm/vlm_engine.py index b8c14b0c08..fe91439f50 100644 --- a/src/megatron/bridge/inference/vlm/vlm_engine.py +++ b/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.""" diff --git a/src/megatron/bridge/models/qwen3_asr/hf_qwen3_asr/modeling_qwen3_asr.py b/src/megatron/bridge/models/qwen3_asr/hf_qwen3_asr/modeling_qwen3_asr.py index 91fe4118ba..21b296a251 100644 --- a/src/megatron/bridge/models/qwen3_asr/hf_qwen3_asr/modeling_qwen3_asr.py +++ b/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**, diff --git a/tests/unit_tests/inference/vlm/test_vlm_engine.py b/tests/unit_tests/inference/vlm/test_vlm_engine.py index 8080eaaf86..81315a1f23 100644 --- a/tests/unit_tests/inference/vlm/test_vlm_engine.py +++ b/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 From 34ab3ea01628942e8c2dd321e3aeaf23187b4422 Mon Sep 17 00:00:00 2001 From: yaoyu-33 Date: Mon, 18 May 2026 14:07:23 -0700 Subject: [PATCH 2/2] [mcore] chore: Remove stale compatibility guards Signed-off-by: yaoyu-33 --- .../bridge/inference/vlm/_mcore_compat.py | 3 +- src/megatron/bridge/models/gpt_provider.py | 7 +-- .../bridge/models/mamba/mamba_provider.py | 50 ++----------------- .../bridge/recipes/utils/optimizer_utils.py | 10 +--- src/megatron/bridge/training/initialize.py | 4 +- src/megatron/bridge/training/optim.py | 18 +------ src/megatron/bridge/training/state.py | 36 +++---------- .../bridge/training/tokenizers/tokenizer.py | 30 ++--------- .../models/mamba/test_mamba_provider.py | 28 ++--------- .../recipes/utils/test_optimizer_utils.py | 4 +- 10 files changed, 25 insertions(+), 165 deletions(-) diff --git a/src/megatron/bridge/inference/vlm/_mcore_compat.py b/src/megatron/bridge/inference/vlm/_mcore_compat.py index eccc00dc61..724378d6ad 100644 --- a/src/megatron/bridge/inference/vlm/_mcore_compat.py +++ b/src/megatron/bridge/inference/vlm/_mcore_compat.py @@ -20,7 +20,8 @@ if "InferenceMode" not in str(exc): raise - # TODO: remove this guard when Megatron-Core dev exposes InferenceMode from megatron.core.inference.utils. + # 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.""" diff --git a/src/megatron/bridge/models/gpt_provider.py b/src/megatron/bridge/models/gpt_provider.py index 9fa191eba9..887b8bc8ac 100644 --- a/src/megatron/bridge/models/gpt_provider.py +++ b/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 diff --git a/src/megatron/bridge/models/mamba/mamba_provider.py b/src/megatron/bridge/models/mamba/mamba_provider.py index fdbcd43c43..f9f2b393a7 100644 --- a/src/megatron/bridge/models/mamba/mamba_provider.py +++ b/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, diff --git a/src/megatron/bridge/recipes/utils/optimizer_utils.py b/src/megatron/bridge/recipes/utils/optimizer_utils.py index df60cfe9fa..61ec16c53e 100644 --- a/src/megatron/bridge/recipes/utils/optimizer_utils.py +++ b/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, diff --git a/src/megatron/bridge/training/initialize.py b/src/megatron/bridge/training/initialize.py index ef0dd88e5c..ba4f381ff2 100644 --- a/src/megatron/bridge/training/initialize.py +++ b/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: diff --git a/src/megatron/bridge/training/optim.py b/src/megatron/bridge/training/optim.py index 658bf85b54..52e2d8a316 100644 --- a/src/megatron/bridge/training/optim.py +++ b/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, diff --git a/src/megatron/bridge/training/state.py b/src/megatron/bridge/training/state.py index 1d21668172..a9452493c1 100644 --- a/src/megatron/bridge/training/state.py +++ b/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 diff --git a/src/megatron/bridge/training/tokenizers/tokenizer.py b/src/megatron/bridge/training/tokenizers/tokenizer.py index 3405bc7e08..dfcd50fe4f 100644 --- a/src/megatron/bridge/training/tokenizers/tokenizer.py +++ b/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 diff --git a/tests/unit_tests/models/mamba/test_mamba_provider.py b/tests/unit_tests/models/mamba/test_mamba_provider.py index 0291de8b78..b442e75132 100644 --- a/tests/unit_tests/models/mamba/test_mamba_provider.py +++ b/tests/unit_tests/models/mamba/test_mamba_provider.py @@ -14,7 +14,6 @@ from unittest.mock import Mock, patch -import pytest import torch from megatron.bridge.models.mamba import mamba_provider @@ -303,37 +302,16 @@ def test_dropout_configuration(self): assert provider.attention_dropout == 0.2 assert provider.layernorm_epsilon == 1e-6 - def test_get_hybrid_total_layer_count_prefers_mcore_helper(self): - """Test helper delegates to MCore when available.""" - mock_counter = Mock(return_value=7) - - with patch.object(mamba_provider, "_mcore_get_hybrid_total_layer_count", mock_counter): - assert mamba_provider._get_hybrid_total_layer_count("M*M*") == 7 - - mock_counter.assert_called_once_with("M*M*") - - def test_get_hybrid_total_layer_count_fallback_supports_pipe_and_mtp(self): - """Test fallback counts only main-decoder layers for newer pattern syntax.""" - with patch.object(mamba_provider, "_mcore_get_hybrid_total_layer_count", None): - assert mamba_provider._get_hybrid_total_layer_count("M-M-|M-M*-/MM/MM") == 9 - - def test_get_hybrid_total_layer_count_fallback_rejects_invalid_symbols(self): - """Test fallback validation matches MCore-style pattern validation.""" - with patch.object(mamba_provider, "_mcore_get_hybrid_total_layer_count", None): - with pytest.raises(ValueError, match="not a valid layer symbol"): - mamba_provider._get_hybrid_total_layer_count("M-A-") - def test_finalize_uses_compatible_hybrid_layer_count(self): - """Test finalize derives num_layers even when older MCore lacks the helper.""" + """Test finalize derives num_layers with MCore's hybrid layer helper.""" provider = MambaModelProvider( hidden_size=768, num_attention_heads=8, hybrid_layer_pattern="M-M-|M-M*-/MM/MM", ) - with patch.object(mamba_provider, "_mcore_get_hybrid_total_layer_count", None): - with patch.object(mamba_provider.TransformerConfig, "finalize", autospec=True) as mock_finalize: - provider.finalize() + with patch.object(mamba_provider.TransformerConfig, "finalize", autospec=True) as mock_finalize: + provider.finalize() assert provider.num_layers == 9 mock_finalize.assert_called_once_with(provider) diff --git a/tests/unit_tests/recipes/utils/test_optimizer_utils.py b/tests/unit_tests/recipes/utils/test_optimizer_utils.py index 9b9892ab62..500201938c 100644 --- a/tests/unit_tests/recipes/utils/test_optimizer_utils.py +++ b/tests/unit_tests/recipes/utils/test_optimizer_utils.py @@ -61,7 +61,7 @@ def test_muon_optimizer_config(self): assert optim_cfg.lr == 3e-4 assert optim_cfg.weight_decay == 0.01 assert optim_cfg.muon_extra_scale_factor == 1.01 - assert getattr(optim_cfg, "muon_use_nesterov", getattr(optim_cfg, "muon_nesterov", None)) is False + assert optim_cfg.muon_nesterov is False assert optim_cfg.muon_momentum == 0.98 assert optim_cfg.bf16 is True @@ -85,7 +85,7 @@ def test_muon_lion_optimizer_config(self): assert optim_cfg.lion_beta2 == 0.95 assert optim_cfg.muon_scalar_optimizer == "lion" assert optim_cfg.muon_extra_scale_factor == 1.01 - assert getattr(optim_cfg, "muon_use_nesterov", getattr(optim_cfg, "muon_nesterov", None)) is False + assert optim_cfg.muon_nesterov is False assert optim_cfg.muon_momentum == 0.98 assert optim_cfg.bf16 is True