From bc20f77d433d29ac3fb830d16c1f610bb8cf314d Mon Sep 17 00:00:00 2001 From: hazelduan Date: Thu, 6 Aug 2026 12:14:40 +0800 Subject: [PATCH 1/3] Fix local attention handling on Ascend --- swift/megatron/model/utils.py | 6 +++--- swift/megatron/pipelines/train/sft.py | 8 +++----- swift/megatron/trainers/utils.py | 2 +- tests/megatron/test_megatron_args.py | 24 ++++++++++++++++++++++++ 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py index 1d01ab933a..ea3a1f2ad8 100644 --- a/swift/megatron/model/utils.py +++ b/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 != 'local') _check_attention_backend(args, config) _check_padding_free(args, config) return config diff --git a/swift/megatron/pipelines/train/sft.py b/swift/megatron/pipelines/train/sft.py index e3fe6a1ffb..af0807ac39 100644 --- a/swift/megatron/pipelines/train/sft.py +++ b/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 != 'local' + 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: diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py index 66d1e5dfd9..d1ea0187c3 100644 --- a/swift/megatron/trainers/utils.py +++ b/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) == 'local': return False return bool(getattr(args, 'use_flash_attn', False)) diff --git a/tests/megatron/test_megatron_args.py b/tests/megatron/test_megatron_args.py index aa36843bb2..26ed84d128 100644 --- a/tests/megatron/test_megatron_args.py +++ b/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,28 @@ 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_local_attention_disables_padding_free(self): + self._skip_if_no_megatron() + from swift.megatron.model.utils import _check_padding_free + + args = SimpleNamespace(padding_free=True) + config = SimpleNamespace(attention_backend=SimpleNamespace(name='local')) + _check_padding_free(args, config) + self.assertFalse(args.padding_free) + + @patch('transformers.utils.is_torch_npu_available', return_value=True) + def test_local_attention_keeps_npu_attention_mask(self, _): + self._skip_if_no_megatron() + from swift.megatron.trainers.utils import _should_use_npu_generated_attention_mask + + args = SimpleNamespace( + task_type='causal_lm', + padding_free=False, + attention_backend=SimpleNamespace(name='local'), + use_flash_attn=True, + ) + self.assertFalse(_should_use_npu_generated_attention_mask(args)) + if __name__ == '__main__': unittest.main() From f08fd2dd4ac76250788ae7d731e12d8d68306599 Mon Sep 17 00:00:00 2001 From: hazelduan Date: Thu, 6 Aug 2026 18:51:50 +0800 Subject: [PATCH 2/3] Support local MCore linears in LoRA targeting --- swift/megatron/utils/utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/swift/megatron/utils/utils.py b/swift/megatron/utils/utils.py index d283cf3ae9..61414c906d 100644 --- a/swift/megatron/utils/utils.py +++ b/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 From e32579f37e14fd53bdc0dc30d938b7b628e75ddb Mon Sep 17 00:00:00 2001 From: hazelduan Date: Thu, 6 Aug 2026 22:10:49 +0800 Subject: [PATCH 3/3] Fix unfused attention selection on Ascend --- swift/megatron/model/utils.py | 2 +- swift/megatron/pipelines/train/sft.py | 2 +- swift/megatron/trainers/utils.py | 2 +- tests/megatron/test_megatron_args.py | 30 +++++++++++++++------------ 4 files changed, 20 insertions(+), 16 deletions(-) diff --git a/swift/megatron/model/utils.py b/swift/megatron/model/utils.py index ea3a1f2ad8..ed55d516dd 100644 --- a/swift/megatron/model/utils.py +++ b/swift/megatron/model/utils.py @@ -73,7 +73,7 @@ def get_mcore_model_config(args, hf_config): kwargs.update(args.megatron_extra_kwargs) config = ModelConfig(**kwargs) if is_torch_npu_available(): - setattr(config, 'use_flash_attn', config.attention_backend.name != 'local') + 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 diff --git a/swift/megatron/pipelines/train/sft.py b/swift/megatron/pipelines/train/sft.py index af0807ac39..e6379cb622 100644 --- a/swift/megatron/pipelines/train/sft.py +++ b/swift/megatron/pipelines/train/sft.py @@ -43,7 +43,7 @@ def __init__(self, args: Optional[Union[List[str], MegatronSftArguments]] = None if apply_mindspeed_patches is not None: megatron_args = asdict(self.args) # MindSpeed requires passing `use_flash_attn` to Megatron to select the attention implementation. - args.use_flash_attn = args.attention_backend.name != 'local' + 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 diff --git a/swift/megatron/trainers/utils.py b/swift/megatron/trainers/utils.py index d1ea0187c3..4dc1f1f4a4 100644 --- a/swift/megatron/trainers/utils.py +++ b/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, 'name', args.attention_backend) == 'local': + if getattr(args.attention_backend, 'name', args.attention_backend) in {'local', 'unfused'}: return False return bool(getattr(args, 'use_flash_attn', False)) diff --git a/tests/megatron/test_megatron_args.py b/tests/megatron/test_megatron_args.py index 26ed84d128..1204d4f261 100644 --- a/tests/megatron/test_megatron_args.py +++ b/tests/megatron/test_megatron_args.py @@ -89,27 +89,31 @@ 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_local_attention_disables_padding_free(self): + def test_non_flash_attention_disables_padding_free(self): self._skip_if_no_megatron() from swift.megatron.model.utils import _check_padding_free - args = SimpleNamespace(padding_free=True) - config = SimpleNamespace(attention_backend=SimpleNamespace(name='local')) - _check_padding_free(args, config) - self.assertFalse(args.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_local_attention_keeps_npu_attention_mask(self, _): + 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 - args = SimpleNamespace( - task_type='causal_lm', - padding_free=False, - attention_backend=SimpleNamespace(name='local'), - use_flash_attn=True, - ) - self.assertFalse(_should_use_npu_generated_attention_mask(args)) + 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__':