[bugfix] fix megatron flash_attn (flash_attention_3) (#5837)

Author: Jintao-Huang

swift/megatron/model/mm_gpt/utils.py

@@ -63,7 +63,8 @@ def __init__(self, config, ignore_init_model_cls=None):
         super().__init__(config)
         args = get_args()
         model_dir = args.model_info.model_dir
-        kwargs = {'attn_impl': 'flash_attn'} if args.attention_backend.name == 'flash' else {}
+        attn_impl = getattr(args, 'attn_impl', None) or 'flash_attn'
+        kwargs = {'attn_impl': attn_impl} if args.attention_backend.name == 'flash' else {}
         ignore_init_model_cls = ignore_init_model_cls or []
         if not isinstance(ignore_init_model_cls, list):
             ignore_init_model_cls = [ignore_init_model_cls]