Model accuracy alignment

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:

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):

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':

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,24 +91,39 @@ 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
 
     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

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