diff --git a/swift/megatron/arguments/megatron_args.py b/swift/megatron/arguments/megatron_args.py index 082ba37ee4..c5aa87556a 100644 --- a/swift/megatron/arguments/megatron_args.py +++ b/swift/megatron/arguments/megatron_args.py @@ -342,6 +342,12 @@ class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin): optimizer_cpu_offload: bool = False optimizer_offload_fraction: float = 1. use_precision_aware_optimizer: bool = False + # Master switch for PaddleFleet<->Megatron bit-alignment patches contributed by + # ningzhengsheng in ms-swift and Megatron-LM. When True the patched (accuracy-aligned) + # code paths are enabled; when False every patched site falls back to the original logic. + # Exported to the `USE_ACCURACY_COMPATIBLE` env var in __post_init__ so both ms-swift and + # Megatron-LM core (incl. standalone/autograd functions without config access) can read it. + use_accuracy_compatible: bool = False main_grads_dtype: Literal['fp32', 'bf16'] = 'fp32' main_params_dtype: Literal['fp32', 'fp16'] = 'fp32' exp_avg_dtype: Literal['fp32', 'fp16', 'bf16', 'fp8'] = 'fp32' @@ -567,6 +573,10 @@ def __post_init__(self): RLHFMegatronArgumentsMixin.__post_init__(self) MegatronTunerMixin.__post_init__(self) os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1') + # Propagate the accuracy-alignment switch to a process-wide env var so that both + # ms-swift and Megatron-LM patched code paths (including moe_utils standalone / + # autograd functions that have no config access) can gate on a single source. + os.environ['USE_ACCURACY_COMPATIBLE'] = '1' if self.use_accuracy_compatible else '0' if self.recompute_granularity == 'none': self.recompute_granularity = None if self.recompute_granularity == 'selective' and self.recompute_method is not None: diff --git a/swift/megatron/model/gpt_bridge.py b/swift/megatron/model/gpt_bridge.py index 7ac9976de0..e919b2091f 100644 --- a/swift/megatron/model/gpt_bridge.py +++ b/swift/megatron/model/gpt_bridge.py @@ -1,6 +1,7 @@ # Copyright (c) ModelScope Contributors. All rights reserved. import math import megatron.core +import os import re import torch import torch.distributed as dist @@ -23,6 +24,15 @@ logger = get_logger() + +def _use_accuracy_compatible() -> bool: + """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. + + Driven by ms-swift's ``use_accuracy_compatible`` arg via the ``USE_ACCURACY_COMPATIBLE`` + env var. Defaults to False so the original TE-folded layernorm layout is used. + """ + return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' + mcore_013 = version.parse(megatron.core.__version__) >= version.parse('0.13.0rc0') EP_PP_SIZE = None @@ -1297,8 +1307,16 @@ def _set_layer_attn(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: boo self._set_state_dict(mg_layer, 'input_layernorm.weight', hf_state_dict, 'input_layernorm.weight', to_mcore) else: hf_state_dict.update(self._set_attn_state(mg_attn, hf_state_dict, 'self_attn.', layer_idx, to_mcore)) - self._set_state_dict(mg_layer, 'self_attention.linear_qkv.layer_norm_weight', hf_state_dict, - 'input_layernorm.weight', to_mcore) + if _use_accuracy_compatible(): + # alignment: with use_transformer_engine=False (local spec), the + # input layernorm weight lives on a standalone `input_layernorm` + # module, not folded into `linear_qkv.layer_norm_weight`. + self._set_state_dict(mg_layer, 'input_layernorm.weight', hf_state_dict, + 'input_layernorm.weight', to_mcore) + else: + # original: TE spec folds the input layernorm into linear_qkv. + self._set_state_dict(mg_layer, 'self_attention.linear_qkv.layer_norm_weight', hf_state_dict, + 'input_layernorm.weight', to_mcore) return hf_state_dict def _set_layer_mlp(self, mg_layer, hf_state_dict, layer_idx: int, to_mcore: bool): diff --git a/swift/megatron/model/model_config.py b/swift/megatron/model/model_config.py index 533518a609..29d5ee8d59 100644 --- a/swift/megatron/model/model_config.py +++ b/swift/megatron/model/model_config.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import os import re import torch.nn.functional as F from dataclasses import dataclass, fields @@ -219,6 +220,11 @@ def _augment_mindspeed_defaults(self): setattr(self, name, value) def __post_init__(self): + # alignment: use_accuracy_compatible switches to the local (non-TE) spec, whose + # standalone LayerNorm modules are incompatible with the persistent LN fusion. + # Disable it here; otherwise keep the Megatron default (persistent LN enabled). + if os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1': + self.persist_layer_norm = False self._augment_mindspeed_defaults() self._format_config() if self.moe_router_dtype.lower() == 'none': diff --git a/swift/megatron/model/register.py b/swift/megatron/model/register.py index 106fdb9a07..2bcfd5d66d 100644 --- a/swift/megatron/model/register.py +++ b/swift/megatron/model/register.py @@ -1,9 +1,12 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import os + import megatron.core from dataclasses import dataclass from megatron.core import mpu from megatron.core.enums import ModelType from megatron.core.models.gpt.gpt_layer_specs import (get_gpt_decoder_block_spec, + get_gpt_layer_local_spec, get_gpt_layer_with_transformer_engine_spec, get_gpt_mtp_block_spec) from packaging import version @@ -24,6 +27,15 @@ logger = get_logger() +def _use_accuracy_compatible() -> bool: + """Runtime switch for the PaddleFleet<->Megatron bit-alignment patches. + + Driven by ms-swift's ``use_accuracy_compatible`` arg via the ``USE_ACCURACY_COMPATIBLE`` + env var. Defaults to False so the original TE-based specs are used when disabled. + """ + return os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' + + @dataclass class MegatronModelMeta: megatron_model_type: str @@ -79,10 +91,13 @@ def __init__(self, args, hf_config): self.model_cls = MultimodalGPTModel if self.args.is_multimodal else GPTModel def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): + # alignment: use_accuracy_compatible disables TE (local spec) to avoid the TE + # DotProductAttention path and match PaddleFleet numerics; otherwise keep TE. + use_te = not _use_accuracy_compatible() if self.config.num_moe_experts: kwargs = {'qk_l2_norm': self.config.qk_l2_norm, 'vp_stage': vp_stage} if self.mcore_013 else {} transformer_layer_spec = get_gpt_decoder_block_spec( - self.config, use_transformer_engine=True, normalization=self.config.normalization, **kwargs) + self.config, use_transformer_engine=use_te, normalization=self.config.normalization, **kwargs) else: transformer_layer_spec = self._get_transformer_layer_spec() return transformer_layer_spec @@ -90,13 +105,25 @@ def get_transformer_layer_spec(self, vp_stage: Optional[int] = None): def _get_transformer_layer_spec(self): config = self.config kwargs = {'qk_l2_norm': config.qk_l2_norm} if self.mcore_013 else {} - transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( - config.num_moe_experts, - self.args.moe_grouped_gemm, - config.qk_layernorm, - config.multi_latent_attention, - **kwargs, - ) + if _use_accuracy_compatible(): + # alignment: use local spec to avoid TE DotProductAttention path, + # which requires a TE backend (FA / cuDNN / unfused) that may be + # unavailable in this environment. + transformer_layer_spec = get_gpt_layer_local_spec( + config.num_moe_experts, + self.args.moe_grouped_gemm, + config.qk_layernorm, + config.multi_latent_attention, + **kwargs, + ) + else: + transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec( + config.num_moe_experts, + self.args.moe_grouped_gemm, + config.qk_layernorm, + config.multi_latent_attention, + **kwargs, + ) return transformer_layer_spec def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = None): @@ -109,7 +136,8 @@ def get_mtp_block_spec(self, transformer_layer_spec, vp_stage: Optional[int] = N transformer_layer_spec_for_mtp = transformer_layer_spec kwargs = {'vp_stage': vp_stage} if self.mcore_013 else {} return get_gpt_mtp_block_spec( - self.config, transformer_layer_spec_for_mtp, use_transformer_engine=True, **kwargs) + self.config, transformer_layer_spec_for_mtp, use_transformer_engine=not _use_accuracy_compatible(), + **kwargs) def _set_shared_expert_gate(self, transformer_layer_spec): if (self.config.use_shared_expert_gate and self.config.num_moe_experts diff --git a/swift/template/templates/utils.py b/swift/template/templates/utils.py index 9a07ec8da5..9f840de3be 100644 --- a/swift/template/templates/utils.py +++ b/swift/template/templates/utils.py @@ -1,4 +1,5 @@ # Copyright (c) ModelScope Contributors. All rights reserved. +import os from dataclasses import dataclass, field from typing import Optional @@ -21,6 +22,11 @@ class ChatmlTemplateMeta(TemplateMeta): @dataclass class EmptyTemplateMeta(TemplateMeta): + # alignment: with use_accuracy_compatible (USE_ACCURACY_COMPATIBLE=1) the dummy template + # adds no suffix token, so raw text -> token_ids matches PaddleFleet bit-for-bit. When the + # switch is off, keep the base TemplateMeta default suffix ([['eos_token_id']]). + suffix: Prompt = field( + default_factory=lambda: [] if os.environ.get('USE_ACCURACY_COMPATIBLE', '0') == '1' else [['eos_token_id']]) prefix: Prompt = field(default_factory=list) prompt: Prompt = field(default_factory=lambda: ['{{QUERY}}']) chat_sep: Optional[Prompt] = None