From 2679e800a3be1cde1f85eca8ebae1e1f79fd31b4 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Fri, 10 Jul 2026 16:12:30 +0800
Subject: [PATCH 1/4] init

---
 swift/megatron/arguments/megatron_args.py |  57 ++++++--
 swift/megatron/init.py                    |   1 -
 swift/megatron/model/utils.py             | 166 ++++++++++++++++++++++
 swift/megatron/trainers/utils.py          |   5 +
 swift/megatron/utils/megatron_lm_utils.py |   2 +-
 swift/pipelines/train/rlhf.py             |   3 +
 swift/rlhf_trainers/gkd_trainer.py        |   2 +-
 swift/rlhf_trainers/grpo_trainer.py       |   2 +-
 swift/rlhf_trainers/rollout_mixin.py      |   3 +
 9 files changed, 229 insertions(+), 12 deletions(-)

diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py
index 42e379502d..270b5723e6 100644
--- a/swift/megatron/arguments/megatron_args.py
+++ b/swift/megatron/arguments/megatron_args.py
@@ -665,7 +665,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
@@ -726,7 +727,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:
@@ -763,6 +764,27 @@ 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".')
+            # Skip mcore-bridge model_meta lookup for megatron-bridge backend.
+            # Model support will be validated via AutoBridge.supports() in get_mcore_model.
+        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.')
@@ -770,6 +792,7 @@ def __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:
@@ -790,9 +813,22 @@ 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':
+            # megatron-bridge does not use mcore-bridge's ModelMeta/ModelConfig.
+            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". Only causal_lm is supported. '
+                    f'Please use bridge_backend="mcore-bridge" for other task types.')
+        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')
@@ -937,9 +973,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:
@@ -1022,6 +1061,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', [])]
diff --git a/swift/megatron/init.py b/swift/megatron/init.py
index c16e2b12f2..d93375da5b 100644
--- a/swift/megatron/init.py
+++ b/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:
diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py
index acc2a7e48b..1b2a02fa48 100644
--- a/swift/megatron/model/utils.py
+++ b/swift/megatron/model/utils.py
@@ -1,5 +1,7 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from dataclasses import fields
+from typing import Any, Generator, Optional, Tuple
+
 from mcore_bridge import ModelConfig
 from mcore_bridge import get_mcore_model as _get_mcore_model
 from mcore_bridge import hf_to_mcore_config
@@ -79,9 +81,173 @@ def get_mcore_model_config(args, hf_config):
     _check_padding_free(args, 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)
+            HfConfigFactory.del_config_attr(hf_config, 'quantization_config')
+            hf_config.save_pretrained(output_dir)
+            if processor is not None:
+                save_checkpoint(
+                    None, processor, output_dir,
+                    model_dirs=[args.model_dir],
+                    additional_saved_files=getattr(args, 'additional_saved_files', None))
+        if dist.is_initialized():
+            dist.barrier()
 
 def get_mcore_model(args, hf_config):
+    bridge_backend = getattr(args, 'bridge_backend', 'mcore-bridge')
+    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
+    if overrides.get('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)
+
+    # --- 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
diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py
index 16ca70d9a2..998a92b41f 100644
--- a/swift/megatron/trainers/utils.py
+++ b/swift/megatron/trainers/utils.py
@@ -399,6 +399,11 @@ 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
+    # padding_free mode does not produce attention_mask in the batch.
+    # mcore-bridge's GPTModel.forward has attention_mask=None default,
+    # but megatron-core's GPTModel.forward requires it as positional arg.
+    # Setting None is safe for both backends.
+    batch.setdefault('attention_mask', None)
     batch = get_batch_on_this_cp_rank(args, batch)
     return batch
 
diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py
index 0484d40107..4e2172e424 100644
--- a/swift/megatron/utils/megatron_lm_utils.py
+++ b/swift/megatron/utils/megatron_lm_utils.py
@@ -670,7 +670,7 @@ def warmup_jit_function(config, args):
     for bias_grad, input_grad in zip([True, True], [False, True]):
         bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad
         for _ in range(5):
-            if config.swiglu:
+            if getattr(config, 'gated_linear_unit', getattr(config, 'swiglu', True)):
                 output = bias_swiglu(input_tensor, bias)
             else:
                 output = bias_gelu(bias, input_tensor)
diff --git a/swift/pipelines/train/rlhf.py b/swift/pipelines/train/rlhf.py
index 1b50d1b8f3..cdb6f3322f 100644
--- a/swift/pipelines/train/rlhf.py
+++ b/swift/pipelines/train/rlhf.py
@@ -239,6 +239,9 @@ def _get_trainer_kwargs(self):
             if self.args.teacher_model_server:
                 trainer_kwargs['teacher_model_server'] = self.args.teacher_model_server
             trainer_kwargs['teacher_use_disable_adapter'] = getattr(self.args, '_teacher_use_disable_adapter', False)
+            # teacher_tag_key lives on the CLI RLHFArguments (TeacherModelArguments), not on the
+            # trainer's GRPOConfig/GKDConfig, so it must ride the trainer-kwargs channel.
+            trainer_kwargs['teacher_tag_key'] = getattr(self.args, 'teacher_tag_key', 'dataset')
         if self.args.rlhf_type == 'gkd':
             trainer_kwargs['gkd_logits_topk'] = self.args.gkd_logits_topk
         return trainer_kwargs
diff --git a/swift/rlhf_trainers/gkd_trainer.py b/swift/rlhf_trainers/gkd_trainer.py
index d6011e72e4..725dea3092 100644
--- a/swift/rlhf_trainers/gkd_trainer.py
+++ b/swift/rlhf_trainers/gkd_trainer.py
@@ -380,7 +380,7 @@ def _fetch_and_assemble_teacher_logprobs(self, chunks):
                 handle, topk=self.gkd_logits_topk, teacher_client=client),
             scatter_fn=self._scatter_teacher_parsed,
             is_main_process=self.accelerator.is_main_process,
-            tag_key=getattr(self.args, 'teacher_tag_key', 'dataset'))
+            tag_key=self.teacher_tag_key)
 
         per_chunk_parsed, offset = [], 0
         for cs in chunk_sizes:
diff --git a/swift/rlhf_trainers/grpo_trainer.py b/swift/rlhf_trainers/grpo_trainer.py
index caada6142c..c077c7bbd7 100644
--- a/swift/rlhf_trainers/grpo_trainer.py
+++ b/swift/rlhf_trainers/grpo_trainer.py
@@ -594,7 +594,7 @@ def _assemble_teacher_api_logps(self, samples: List[GRPOSample], batch_encoded_i
             infer_fn=lambda handle, client: self._infer_teacher_requests(handle, topk=0, teacher_client=client),
             scatter_fn=self._scatter_teacher_parsed,
             is_main_process=self.accelerator.is_main_process,
-            tag_key=getattr(self.args, 'teacher_tag_key', 'dataset'))
+            tag_key=self.teacher_tag_key)
 
         offset = 0
         for batch_encoded, chunk in zip(batch_encoded_inputs, chunk_samples_list):
diff --git a/swift/rlhf_trainers/rollout_mixin.py b/swift/rlhf_trainers/rollout_mixin.py
index 87e6901691..160a2eb8fc 100644
--- a/swift/rlhf_trainers/rollout_mixin.py
+++ b/swift/rlhf_trainers/rollout_mixin.py
@@ -189,6 +189,9 @@ def _pop_teacher_kwargs(self, kwargs: Dict[str, Any]) -> None:
         self._teacher_deepspeed_config = kwargs.pop('teacher_deepspeed_config', None)
         self.teacher_model_server = kwargs.pop('teacher_model_server', None)
         self._teacher_use_disable_adapter = kwargs.pop('teacher_use_disable_adapter', False)
+        # From CLI RLHFArguments via trainer kwargs: the training config (self.args) has no such
+        # field, so reading getattr(self.args, 'teacher_tag_key') would silently yield 'dataset'.
+        self.teacher_tag_key = kwargs.pop('teacher_tag_key', 'dataset')
 
     def _has_teacher_explicit(self) -> bool:
         return (getattr(self, '_teacher_model', None) is not None

From 181602539836fd91cf58d2839147f277fc28674a Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Fri, 10 Jul 2026 16:27:56 +0800
Subject: [PATCH 2/4] update

---
 swift/megatron/arguments/megatron_args.py | 2 --
 swift/megatron/model/utils.py             | 2 +-
 swift/megatron/trainers/utils.py          | 4 ----
 swift/megatron/utils/megatron_lm_utils.py | 2 +-
 swift/pipelines/train/rlhf.py             | 3 ---
 swift/rlhf_trainers/rollout_mixin.py      | 3 ---
 6 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py
index c43588e81e..195ba7d11f 100644
--- a/swift/megatron/arguments/megatron_args.py
+++ b/swift/megatron/arguments/megatron_args.py
@@ -779,8 +779,6 @@ def _check_bridge_backend(self):
                 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".')
-            # Skip mcore-bridge model_meta lookup for megatron-bridge backend.
-            # Model support will be validated via AutoBridge.supports() in get_mcore_model.
         else:
             require_version('mcore-bridge>=1.4.0',
                             'Please install mcore-bridge via `pip install mcore-bridge -U`')
diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py
index 1b2a02fa48..906396c20a 100644
--- a/swift/megatron/model/utils.py
+++ b/swift/megatron/model/utils.py
@@ -159,7 +159,7 @@ def save_weights(self, models, output_dir, peft_format=False, max_shard_size='5G
             dist.barrier()
 
 def get_mcore_model(args, hf_config):
-    bridge_backend = getattr(args, 'bridge_backend', 'mcore-bridge')
+    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)
diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py
index 89cb71507a..66d1e5dfd9 100644
--- a/swift/megatron/trainers/utils.py
+++ b/swift/megatron/trainers/utils.py
@@ -399,10 +399,6 @@ 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
-    # padding_free mode does not produce attention_mask in the batch.
-    # mcore-bridge's GPTModel.forward has attention_mask=None default,
-    # but megatron-core's GPTModel.forward requires it as positional arg.
-    # Setting None is safe for both backends.
     batch.setdefault('attention_mask', None)
     batch = get_batch_on_this_cp_rank(args, batch)
     return batch
diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py
index 37a22cdea8..72fec0c49b 100644
--- a/swift/megatron/utils/megatron_lm_utils.py
+++ b/swift/megatron/utils/megatron_lm_utils.py
@@ -670,7 +670,7 @@ def warmup_jit_function(config, args):
     for bias_grad, input_grad in zip([True, True], [False, True]):
         bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad
         for _ in range(5):
-            if getattr(config, 'gated_linear_unit', getattr(config, 'swiglu', True)):
+            if getattr(config, 'swiglu', getattr(config, 'gated_linear_unit', True)):
                 output = bias_swiglu(input_tensor, bias)
             else:
                 output = bias_gelu(bias, input_tensor)
diff --git a/swift/pipelines/train/rlhf.py b/swift/pipelines/train/rlhf.py
index cdb6f3322f..1b50d1b8f3 100644
--- a/swift/pipelines/train/rlhf.py
+++ b/swift/pipelines/train/rlhf.py
@@ -239,9 +239,6 @@ def _get_trainer_kwargs(self):
             if self.args.teacher_model_server:
                 trainer_kwargs['teacher_model_server'] = self.args.teacher_model_server
             trainer_kwargs['teacher_use_disable_adapter'] = getattr(self.args, '_teacher_use_disable_adapter', False)
-            # teacher_tag_key lives on the CLI RLHFArguments (TeacherModelArguments), not on the
-            # trainer's GRPOConfig/GKDConfig, so it must ride the trainer-kwargs channel.
-            trainer_kwargs['teacher_tag_key'] = getattr(self.args, 'teacher_tag_key', 'dataset')
         if self.args.rlhf_type == 'gkd':
             trainer_kwargs['gkd_logits_topk'] = self.args.gkd_logits_topk
         return trainer_kwargs
diff --git a/swift/rlhf_trainers/rollout_mixin.py b/swift/rlhf_trainers/rollout_mixin.py
index 31c3584202..babe81abe4 100644
--- a/swift/rlhf_trainers/rollout_mixin.py
+++ b/swift/rlhf_trainers/rollout_mixin.py
@@ -189,9 +189,6 @@ def _pop_teacher_kwargs(self, kwargs: Dict[str, Any]) -> None:
         self._teacher_deepspeed_config = kwargs.pop('teacher_deepspeed_config', None)
         self.teacher_model_server = kwargs.pop('teacher_model_server', None)
         self._teacher_use_disable_adapter = kwargs.pop('teacher_use_disable_adapter', False)
-        # From CLI RLHFArguments via trainer kwargs: the training config (self.args) has no such
-        # field, so reading getattr(self.args, 'teacher_tag_key') would silently yield 'dataset'.
-        self.teacher_tag_key = kwargs.pop('teacher_tag_key', 'dataset')
 
     def _has_teacher_explicit(self) -> bool:
         return (getattr(self, '_teacher_model', None) is not None

From d66b049ea1e2f430ea65a051357d4e2179fd5ee9 Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Fri, 10 Jul 2026 17:59:26 +0800
Subject: [PATCH 3/4] update

---
 swift/megatron/arguments/megatron_args.py | 25 +++----
 swift/megatron/model/utils.py             | 91 ++++++++++++++++-------
 2 files changed, 74 insertions(+), 42 deletions(-)

diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py
index 195ba7d11f..fa475a3627 100644
--- a/swift/megatron/arguments/megatron_args.py
+++ b/swift/megatron/arguments/megatron_args.py
@@ -772,16 +772,13 @@ def _check_bridge_backend(self):
             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".')
+                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".')
+                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`')
+            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()
 
@@ -814,17 +811,13 @@ def __post_init__(self):
         self.model_dir = self.model_info.model_dir
         self.is_multimodal = self.model_meta.is_multimodal
         if self.bridge_backend == 'megatron-bridge':
-            # megatron-bridge does not use mcore-bridge's ModelMeta/ModelConfig.
             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.')
+                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". Only causal_lm is supported. '
-                    f'Please use bridge_backend="mcore-bridge" for other task types.')
+                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:
diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py
index 906396c20a..8e0853493d 100644
--- a/swift/megatron/model/utils.py
+++ b/swift/megatron/model/utils.py
@@ -1,11 +1,10 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 from dataclasses import fields
-from typing import Any, Generator, Optional, Tuple
-
 from mcore_bridge import ModelConfig
 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
 
@@ -81,6 +80,7 @@ def get_mcore_model_config(args, hf_config):
     _check_padding_free(args, config)
     return config
 
+
 class MegatronBridgeBackend:
     """Adapter for NVIDIA ``megatron.bridge.AutoBridge``.
 
@@ -99,24 +99,27 @@ 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):
+    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.')
+            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]:
+    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.')
+            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
@@ -128,12 +131,16 @@ def export_weights(self, models, target_device=None, only_master_rank=False,
                 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:
+    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.')
+            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)
@@ -145,19 +152,49 @@ def save_weights(self, models, output_dir, peft_format=False, max_shard_size='5G
         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,
+                    None,
+                    processor,
+                    output_dir,
                     model_dirs=[args.model_dir],
-                    additional_saved_files=getattr(args, 'additional_saved_files', None))
+                    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':
@@ -177,10 +214,9 @@ def _get_megatron_bridge_model(args, hf_config):
     # 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.')
+        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)
@@ -211,7 +247,10 @@ def _get_megatron_bridge_model(args, hf_config):
     dtype = getattr(args, 'torch_dtype', None)
 
     # MoE: if no experts, force EP/ETP to 1
-    if overrides.get('num_moe_experts') is None:
+    # 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
 

From 30bb30b59395ad382aac4ec9b357021604f4dbed Mon Sep 17 00:00:00 2001
From: hjh0119 <hujinghan.hjh@alibaba-inc.com>
Date: Fri, 10 Jul 2026 20:54:48 +0800
Subject: [PATCH 4/4] fix swiglu

---
 swift/megatron/model/utils.py             | 3 +++
 swift/megatron/utils/megatron_lm_utils.py | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py
index 8e0853493d..0478804fd9 100644
--- a/swift/megatron/model/utils.py
+++ b/swift/megatron/model/utils.py
@@ -274,6 +274,9 @@ def _get_megatron_bridge_model(args, hf_config):
     # --- 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(
diff --git a/swift/megatron/utils/megatron_lm_utils.py b/swift/megatron/utils/megatron_lm_utils.py
index 72fec0c49b..0592110a3d 100644
--- a/swift/megatron/utils/megatron_lm_utils.py
+++ b/swift/megatron/utils/megatron_lm_utils.py
@@ -670,7 +670,7 @@ def warmup_jit_function(config, args):
     for bias_grad, input_grad in zip([True, True], [False, True]):
         bias.requires_grad, input_tensor.requires_grad = bias_grad, input_grad
         for _ in range(5):
-            if getattr(config, 'swiglu', getattr(config, 'gated_linear_unit', True)):
+            if config.swiglu:
                 output = bias_swiglu(input_tensor, bias)
             else:
                 output = bias_gelu(bias, input_tensor)