# Copyright (c) ModelScope Contributors. All rights reserved. from dataclasses import fields 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 logger = get_logger() def _check_attention_backend(args, config): """Validate attention backend compatibility with configuration.""" attention_backend = config.attention_backend.name if attention_backend == 'flash' and config.softmax_type == 'learnable': raise ValueError(f'Attention backend "{attention_backend}" does not support learnable softmax_type.') def _check_padding_free(args, config): """Validate and adjust padding_free setting based on configuration constraints.""" if not args.padding_free: return attention_backend = config.attention_backend.name message = None if config.experimental_attention_variant == 'dsa': message = 'DSA is not supported in padding-free mode' elif attention_backend == 'unfused': message = f'Attention backend "{attention_backend}" is not supported in padding-free mode' if message: logger.warning(f'{message}. Setting args.padding_free to False.') args.padding_free = False def get_mcore_model_config(args, hf_config): kwargs = hf_to_mcore_config(hf_config) kwargs['mcore_model_type'] = args.megatron_model_meta.model_type kwargs['hf_config'] = hf_config for f in fields(ModelConfig): key, value = f.name, getattr(args, f.name, None) if value is None or isinstance(value, (list, tuple)) and len(value) == 0: continue kwargs[key] = value if args.task_type == 'seq_cls': args.problem_type = args.problem_type or getattr(hf_config, 'problem_type', None) logger.info(f'args.problem_type: {args.problem_type}') kwargs['params_dtype'] = args.torch_dtype kwargs['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers kwargs['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers kwargs['fp4_param'] = args.fp4_param_gather kwargs['fp8_param'] = args.fp8_param_gather swiglu = kwargs.get('swiglu', True) add_bias_linear = kwargs.get('add_bias_linear', False) num_moe_experts = kwargs.get('num_moe_experts', None) position_embedding_type = kwargs.get('position_embedding_type', 'rope') if position_embedding_type != 'rope': kwargs['apply_rope_fusion'] = False if not swiglu and not add_bias_linear: kwargs['bias_activation_fusion'] = False if add_bias_linear and num_moe_experts and args.moe_grouped_gemm: kwargs['bias_dropout_fusion'] = False if num_moe_experts is None: kwargs['expert_model_parallel_size'] = 1 kwargs['expert_tensor_parallel_size'] = 1 if args.router_replay_mode != 'disabled': kwargs['moe_enable_routing_replay'] = True if args.megatron_extra_kwargs: kwargs.update(args.megatron_extra_kwargs) config = ModelConfig(**kwargs) if is_torch_npu_available() and getattr(args, 'attention_backend', 'flash') != 'local': setattr(config, 'use_flash_attn', True) _check_attention_backend(args, 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) 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