diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py index 2e98ad3326..fa475a3627 100644 --- a/swift/megatron/arguments/megatron_args.py +++ b/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,6 +766,22 @@ 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.') @@ -772,6 +789,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: @@ -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', [])] 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..0478804fd9 100644 --- a/swift/megatron/model/utils.py +++ b/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 diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py index de899317ca..66d1e5dfd9 100644 --- a/swift/megatron/trainers/utils.py +++ b/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