support optional megatron-bridge backend

swift/megatron/arguments/megatron_args.py

@@ -667,7 +667,8 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin):
     mtp_decoder_input_detach: bool = False
     mtp_shared_weights: bool = False
 
-    # mcore-bridge
+    # mcore-bridge / megatron-bridge
+    bridge_backend: Literal['mcore-bridge', 'megatron-bridge'] = 'mcore-bridge'
     model: Optional[str] = None
     model_type: Optional[str] = None
     save_safetensors: bool = True
@@ -728,7 +729,7 @@ def load_args_config(ckpt_dir: Optional[str]) -> Dict[str, Any]:
             with open(args_path, 'r', encoding='utf-8') as f:
                 old_args = json.load(f)
             keys = list(f.name for f in fields(MegatronTunerMixin))
-            keys += ['mcore_model', 'task_type', 'num_labels']
+            keys += ['mcore_model', 'task_type', 'num_labels', 'bridge_backend']
             for key in keys:
                 old_value = old_args.get(key)
                 if old_value is not None:
@@ -765,13 +766,30 @@ def _check_mcore_bridge(self):
                 raise ValueError('`tuner_type="lora_llm"` is not supported when `language_model_only=True`. '
                                  'Please use `tuner_type="lora"` instead.')
 
+    def _check_bridge_backend(self):
+        """Validate bridge_backend and associated constraints."""
+        if self.bridge_backend == 'megatron-bridge':
+            try:
+                import megatron.bridge
+            except ImportError:
+                raise ImportError('bridge_backend="megatron-bridge" requires the `megatron-bridge` package. '
+                                  'Install it via `pip install megatron-bridge` or use bridge_backend="mcore-bridge".')
+            if self.tuner_type != 'full':
+                raise ValueError('LoRA training is not yet supported with bridge_backend="megatron-bridge". '
+                                 'Please use bridge_backend="mcore-bridge" for LoRA, or set tuner_type="full".')
+        else:
+            require_version('mcore-bridge>=1.4.0', 'Please install mcore-bridge via `pip install mcore-bridge -U`')
+            from swift.megatron.init import _patch_mcore_bridge
+            _patch_mcore_bridge()
+
     def __post_init__(self):
         if self.tuner_type != 'full':
             require_version('peft>=0.15', 'Please install peft>=0.15 to use LoRA in Megatron-SWIFT.')
         RLHFMegatronArgumentsMixin.__post_init__(self)
         MegatronTunerMixin.__post_init__(self)
         os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1')
         self._check_mcore_bridge()
+        self._check_bridge_backend()
         if self.recompute_granularity == 'none':
             self.recompute_granularity = None
         if self.recompute_granularity == 'selective' and self.recompute_method is not None:
@@ -792,9 +810,18 @@ def __post_init__(self):
         self.model_type = self.model_info.model_type
         self.model_dir = self.model_info.model_dir
         self.is_multimodal = self.model_meta.is_multimodal
-        self.megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.model_meta))
-        if self.megatron_model_meta is None:
-            raise ValueError(f'Model: {self.model} is not supported.')
+        if self.bridge_backend == 'megatron-bridge':
+            self.megatron_model_meta = None
+            if self.is_multimodal:
+                raise ValueError('Multimodal training is not yet supported with bridge_backend="megatron-bridge". '
+                                 'Please use bridge_backend="mcore-bridge" for multimodal models.')
+            if self.task_type not in (None, 'causal_lm'):
+                raise ValueError(f'task_type={self.task_type!r} is not yet supported with '
+                                 f'bridge_backend="megatron-bridge".')
+        else:
+            self.megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.model_meta))
+            if self.megatron_model_meta is None:
+                raise ValueError(f'Model: {self.model} is not supported.')
         self._init_teacher_model()
         if self.apply_wd_to_qk_layernorm and self.model_type not in {'qwen3_next', 'qwen3_5', 'qwen3_5_moe'}:
             raise ValueError('apply_wd_to_qk_layernorm is only supported for qwen3_next, qwen3_5 and qwen3_5_moe')
@@ -939,9 +966,12 @@ def _init_teacher_model(self):
             self.teacher_model, model_type=self.teacher_model_type, use_hf=self.use_hf, hub_token=self.hub_token)
         self.teacher_model_type = self.teacher_model_info.model_type
         self.teacher_model_dir = self.teacher_model_info.model_dir
-        self.teacher_megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.teacher_model_meta))
-        if self.teacher_megatron_model_meta is None:
-            raise ValueError(f'Model: {self.teacher_model} is not supported.')
+        if self.bridge_backend == 'megatron-bridge':
+            self.teacher_megatron_model_meta = None
+        else:
+            self.teacher_megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.teacher_model_meta))
+            if self.teacher_megatron_model_meta is None:
+                raise ValueError(f'Model: {self.teacher_model} is not supported.')
 
     def _init_vpp_size(self):
         if self.pipeline_model_parallel_layout is not None:
@@ -1024,6 +1054,8 @@ def init_iters(self, train_dataset, val_dataset):
             self.eval_iters = 0
 
     def _init_multimodal_full(self):
+        if not self.is_multimodal:
+            return
         visual_cls = self.megatron_model_meta.visual_cls
         if self.tuner_type == 'full' and self.is_multimodal and visual_cls is not None and not self.language_model_only:
             vision_tower = [f'visual.{vit}' for vit in getattr(visual_cls, '_vision_tower', [])]

swift/megatron/init.py

@@ -209,7 +209,6 @@ def init_megatron_env():
     os.environ.pop('VLLM_USE_MODELSCOPE', None)
     logging_level = logging.root.level
     _patch_unified_memory()
-    _patch_mcore_bridge()
     _patch__batched_p2p_ops()
     logging.root.setLevel(logging_level)  # revert logger level
     try:

swift/megatron/model/utils.py

@@ -4,6 +4,7 @@
 from mcore_bridge import get_mcore_model as _get_mcore_model
 from mcore_bridge import hf_to_mcore_config
 from transformers.utils import is_torch_npu_available
+from typing import Any, Generator, Optional, Tuple
 
 from swift.utils import get_logger
 
@@ -80,8 +81,215 @@ def get_mcore_model_config(args, hf_config):
     return config
 
 
+class MegatronBridgeBackend:
+    """Adapter for NVIDIA ``megatron.bridge.AutoBridge``.
+
+    Limitations:
+    - LoRA / PEFT loading is not yet supported.
+    - MLLM is not yet supported
+    - FP8 export is not yet supported.
+    """
+
+    def __init__(self, auto_bridge: Any, hf_config: Optional[Any] = None):
+        self._bridge = auto_bridge
+        self._hf_config = hf_config
+
+    @classmethod
+    def from_hf_config(cls, hf_config) -> 'MegatronBridgeBackend':
+        from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+        return cls(AutoBridge.from_hf_config(hf_config), hf_config)
+
+    def load_weights(self, models, hf_model_dir, peft_format=False, adapter_name='default', converter=None):
+        if peft_format:
+            raise NotImplementedError('LoRA loading via megatron-bridge backend is not yet supported. '
+                                      'Please use bridge_backend="mcore-bridge" for LoRA training.')
+        if converter is not None:
+            logger.warning('converter is not supported by megatron-bridge backend, ignoring.')
+        self._bridge.load_hf_weights(models, hf_path=hf_model_dir)
+
+    def export_weights(self,
+                       models,
+                       target_device=None,
+                       only_master_rank=False,
+                       peft_format=False,
+                       adapter_name='default',
+                       converter=None,
+                       tqdm_desc='Exporting: ',
+                       disable_tqdm=True,
+                       _is_saving=False) -> Generator[Tuple[str, 'torch.Tensor'], None, None]:
+        if peft_format:
+            raise NotImplementedError('LoRA export via megatron-bridge backend is not yet supported. '
+                                      'Please use bridge_backend="mcore-bridge" for LoRA training.')
+        cpu = (target_device == 'cpu')
+        for weight_tuple in self._bridge.export_hf_weights(models, cpu=cpu):
+            key = weight_tuple.param_name
+            tensor = weight_tuple.weight
+            if converter is not None and tensor is not None:
+                kv = converter(key, tensor)
+                if kv is None:
+                    continue
+                key, tensor = kv
+            yield key, tensor
+
+    def save_weights(self,
+                     models,
+                     output_dir,
+                     peft_format=False,
+                     max_shard_size='5GB',
+                     args=None,
+                     processor=None) -> None:
+        if peft_format:
+            raise NotImplementedError('LoRA saving via megatron-bridge backend is not yet supported. '
+                                      'Please use bridge_backend="mcore-bridge" for LoRA training.')
+
+        # 1. Save weights via megatron-bridge (safetensors format)
+        self._bridge.save_hf_weights(models, path=output_dir)
+
+        # 2. Save HF config and tokenizer on rank 0.
+        # We use the original HF config (not the one reconstructed by megatron-bridge)
+        # because the bridge's config-only path may drop fields like num_attention_heads.
+        import torch.distributed as dist
+        is_master = (not dist.is_initialized()) or dist.get_rank() == 0
+        if is_master and args is not None and self._hf_config is not None:
+            from copy import deepcopy
+
+            from swift.model import save_checkpoint
+            from swift.utils import HfConfigFactory
+            hf_config = deepcopy(self._hf_config)
+            llm_config = HfConfigFactory.get_text_config(hf_config)
+
+            # MTP: write back num_nextn_predict_layers
+            mtp_num_layers = getattr(args, 'mtp_num_layers', None)
+            if mtp_num_layers:
+                for key in ['num_nextn_predict_layers', 'mtp_num_hidden_layers']:
+                    if hasattr(llm_config, key):
+                        setattr(llm_config, key, mtp_num_layers)
+                        break
+                else:
+                    llm_config.num_nextn_predict_layers = mtp_num_layers
+
+            HfConfigFactory.del_config_attr(hf_config, 'quantization_config')
+
+            # FP8: write back quantization_config
+            expert_dtype = None
+            fp8_format = getattr(args, 'fp8_format', None)
+            fp8_recipe = getattr(args, 'fp8_recipe', 'delayed')
+            fp8_param = getattr(args, 'fp8_param_gather', False)
+            if fp8_format is not None and fp8_recipe == 'blockwise' and fp8_param:
+                from transformers.utils.quantization_config import FineGrainedFP8Config
+                hf_config.quantization_config = FineGrainedFP8Config()
+                expert_dtype = 'fp8'
+            if getattr(args, 'model_type', None) == 'deepseek_v4':
+                HfConfigFactory.set_config_attr(hf_config, 'expert_dtype', expert_dtype)
+
+            hf_config.save_pretrained(output_dir)
+            if processor is not None:
+                additional_saved_files = getattr(getattr(processor, 'model_meta', None), 'additional_saved_files', None)
+                save_checkpoint(
+                    None,
+                    processor,
+                    output_dir,
+                    model_dirs=[args.model_dir],
+                    additional_saved_files=additional_saved_files)
+        if dist.is_initialized():
+            dist.barrier()
+
+
 def get_mcore_model(args, hf_config):
+    bridge_backend = args.bridge_backend
+    if bridge_backend == 'megatron-bridge':
+        return _get_megatron_bridge_model(args, hf_config)
     config = get_mcore_model_config(args, hf_config)
     models = _get_mcore_model(config)
 
     return models
+
+
+def _get_megatron_bridge_model(args, hf_config):
+    import dataclasses
+
+    backend = MegatronBridgeBackend.from_hf_config(hf_config)
+    auto_bridge = backend._bridge
+
+    # Validate model support via AutoBridge.supports()
+    from megatron.bridge.models.conversion.auto_bridge import AutoBridge
+    if not AutoBridge.supports(hf_config):
+        raise ValueError(f'Model {getattr(hf_config, "model_type", "unknown")} is not supported by '
+                         f'megatron-bridge. Please use bridge_backend="mcore-bridge" or check '
+                         f'AutoBridge.list_supported_models() for supported architectures.')
+
+    # --- Step 1: Get provider (GPTModelProvider, which extends TransformerConfig) ---
+    provider = auto_bridge.to_megatron_provider(load_weights=False)
+
+    # --- Step 2: Build overrides from args ---
+    # Auto-match: iterate over provider's dataclass fields and pick up matching args fields.
+    # This mirrors mcore-bridge's get_mcore_model_config which does:
+    #   for f in fields(ModelConfig): kwargs[f.name] = getattr(args, f.name, None)
+    overrides = {}
+    provider_fields = {f.name for f in dataclasses.fields(provider)}
+    for field_name in provider_fields:
+        value = getattr(args, field_name, None)
+        if value is None or (isinstance(value, (list, tuple)) and len(value) == 0):
+            continue
+        overrides[field_name] = value
+
+    # Explicit field name mappings (args name → provider field name)
+    explicit_mappings = {
+        'decoder_first_pipeline_num_layers': 'num_layers_in_first_pipeline_stage',
+        'decoder_last_pipeline_num_layers': 'num_layers_in_last_pipeline_stage',
+    }
+    for args_key, provider_key in explicit_mappings.items():
+        value = getattr(args, args_key, None)
+        if value is not None and provider_key in provider_fields:
+            overrides[provider_key] = value
+
+    # dtype
+    dtype = getattr(args, 'torch_dtype', None)
+
+    # MoE: if no experts, force EP/ETP to 1
+    # num_moe_experts comes from HF config (parsed by AutoBridge into provider),
+    # not from args — so check provider too.
+    num_moe_experts = overrides.get('num_moe_experts') or getattr(provider, 'num_moe_experts', None)
+    if num_moe_experts is None:
+        overrides['expert_model_parallel_size'] = 1
+        overrides['expert_tensor_parallel_size'] = 1
+
+    # Router replay
+    if getattr(args, 'router_replay_mode', 'disabled') != 'disabled':
+        if 'moe_enable_routing_replay' in provider_fields:
+            overrides['moe_enable_routing_replay'] = True
+
+    # megatron_extra_kwargs (user-specified raw overrides)
+    if getattr(args, 'megatron_extra_kwargs', None):
+        overrides.update(args.megatron_extra_kwargs)
+
+    # padding_free requires variable_seq_lengths=True so that RotaryEmbedding
+    # generates freqs matching the actual packed sequence length (cu_seqlens[-1])
+    # instead of the fixed seq_length. Without this, mcore-bridge's patcher
+    # use_batched_rope check fails and falls back to the original
+    # _apply_rotary_pos_emb_thd which calls torch.split on a padded tensor.
+    if getattr(args, 'padding_free', False) and 'variable_seq_lengths' in provider_fields:
+        overrides['variable_seq_lengths'] = True
+
+    # --- Step 3: Apply overrides and finalize ---
+    provider.apply_overrides_and_finalize(dtype=dtype, overrides=overrides)
+
+    import torch.nn.functional as F
+    provider.swiglu = (provider.gated_linear_unit and provider.activation_func is F.silu)
+
+    # --- Step 4: Create raw models (no DDP/Float16 wrapping) ---
+    # swift's wrap_model handles DDP/Float16 wrapping with the correct DDP config from args.
+    models = provider.provide_distributed_model(
+        wrap_with_ddp=False,
+        mixed_precision_wrapper=None,
+        use_cpu_initialization=getattr(args, 'use_cpu_initialization', False),
+    )
+    if not isinstance(models, list):
+        models = [models]
+
+    # --- Step 5: Attach backend to model.config.bridge ---
+    for model in models:
+        model.config.bridge = backend
+
+    logger.info('Created Megatron model via megatron-bridge backend')
+    return models

swift/megatron/trainers/utils.py

@@ -399,6 +399,7 @@ def prepare_batch(args, data, vp_stage=None):
             batch['packed_seq_params'].seq_lens = torch.tensor(seq_lens, device=text_position_ids.device)
         if num_samples is not None:
             batch['packed_seq_params'].num_samples = num_samples
+    batch.setdefault('attention_mask', None)
     batch = get_batch_on_this_cp_rank(args, batch)
     return batch