Fix local/unfused Megatron attention and LoRA on Ascend

swift/megatron/model/utils.py

@@ -26,7 +26,7 @@ def _check_padding_free(args, config):
     attention_backend = config.attention_backend.name
     message = None
 
-    if attention_backend == 'unfused':
+    if attention_backend in {'local', 'unfused'}:
         message = f'Attention backend "{attention_backend}" is not supported in padding-free mode'
 
     if message:
@@ -72,8 +72,8 @@ def get_mcore_model_config(args, hf_config):
     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)
+    if is_torch_npu_available():
+        setattr(config, 'use_flash_attn', config.attention_backend.name not in {'local', 'unfused'})
     _check_attention_backend(args, config)
     _check_padding_free(args, config)
     return config

swift/megatron/pipelines/train/sft.py

@@ -42,11 +42,9 @@ def __init__(self, args: Optional[Union[List[str], MegatronSftArguments]] = None
         args = self.args
         if apply_mindspeed_patches is not None:
             megatron_args = asdict(self.args)
-            if args.attention_backend != 'local':
-                # MindSpeed requires passing `use_flash_attn` to Megatron
-                # to enable flash attention on Ascend NPU.
-                args.use_flash_attn = True
-                megatron_args['use_flash_attn'] = True
+            # MindSpeed requires passing `use_flash_attn` to Megatron to select the attention implementation.
+            args.use_flash_attn = args.attention_backend.name not in {'local', 'unfused'}
+            megatron_args['use_flash_attn'] = args.use_flash_attn
             apply_mindspeed_patches(megatron_args)
         template_cls = args.template_meta.template_cls
         if args.model_meta.is_multimodal and template_cls and template_cls.use_model:

swift/megatron/trainers/utils.py

@@ -356,7 +356,7 @@ def _should_use_npu_generated_attention_mask(args) -> bool:
         return False
     if args.task_type != 'causal_lm' or args.padding_free:
         return False
-    if getattr(args, 'attention_backend', None) == 'local':
+    if getattr(args.attention_backend, 'name', args.attention_backend) in {'local', 'unfused'}:
         return False
     return bool(getattr(args, 'use_flash_attn', False))
 

swift/megatron/utils/utils.py

@@ -8,6 +8,7 @@
 from megatron.core.models.common.embeddings.language_model_embedding import LanguageModelEmbedding
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing
+from megatron.core.tensor_parallel.layers import ColumnParallelLinear, RowParallelLinear
 from megatron.core.transformer.moe.router import TopKRouter
 from torch import nn
 from transformers.utils import is_torch_npu_available
@@ -24,7 +25,8 @@ def find_all_linears(model, extra_layers=None):
 
     def _cond(name, module):
         if (extra_layers and isinstance(module, tuple(extra_layers))) or name != 'output_layer' and isinstance(
-                module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, nn.Linear)):
+                module, (TELinear, TELayerNormColumnParallelLinear, TEGroupedLinear, ColumnParallelLinear,
+                         RowParallelLinear, nn.Linear)):
             return True
         return False
 

tests/megatron/test_megatron_args.py

@@ -1,4 +1,6 @@
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 
 class TestMegatronArgs(unittest.TestCase):
@@ -87,6 +89,32 @@ def test_megatron_base_args_fields(self):
         for field_name in expected_fields:
             self.assertIn(field_name, field_names, f'MegatronArguments missing field: {field_name}')
 
+    def test_non_flash_attention_disables_padding_free(self):
+        self._skip_if_no_megatron()
+        from swift.megatron.model.utils import _check_padding_free
+
+        for attention_backend in ('local', 'unfused'):
+            with self.subTest(attention_backend=attention_backend):
+                args = SimpleNamespace(padding_free=True)
+                config = SimpleNamespace(attention_backend=SimpleNamespace(name=attention_backend))
+                _check_padding_free(args, config)
+                self.assertFalse(args.padding_free)
+
+    @patch('transformers.utils.is_torch_npu_available', return_value=True)
+    def test_non_flash_attention_keeps_npu_attention_mask(self, _):
+        self._skip_if_no_megatron()
+        from swift.megatron.trainers.utils import _should_use_npu_generated_attention_mask
+
+        for attention_backend in ('local', 'unfused'):
+            with self.subTest(attention_backend=attention_backend):
+                args = SimpleNamespace(
+                    task_type='causal_lm',
+                    padding_free=False,
+                    attention_backend=SimpleNamespace(name=attention_backend),
+                    use_flash_attn=True,
+                )
+                self.assertFalse(_should_use_npu_generated_attention_mask(args))
+
 
 if __name__ == '__main__':
     unittest.main()