From 4b1d77e6e49f215227d2bc847f791e9db11eafa5 Mon Sep 17 00:00:00 2001 From: Jintao Huang Date: Thu, 21 May 2026 16:33:57 +0800 Subject: [PATCH] test_convert_precision suppot attention_fused --- swift/megatron/utils/convert_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/swift/megatron/utils/convert_utils.py b/swift/megatron/utils/convert_utils.py index 279e1ee3a5..d46b047e02 100644 --- a/swift/megatron/utils/convert_utils.py +++ b/swift/megatron/utils/convert_utils.py @@ -175,7 +175,7 @@ def broadcast_mg_logits(mg_logits=None, src_rank=None): @contextmanager -def _patch_attention_flash_fp32(compute_dtype): +def _patch_attention_fp32(compute_dtype): forward = TEDotProductAttention.forward def new_forward(self, query_layer, key_layer, value_layer, *args, **kwargs): @@ -252,11 +252,11 @@ def test_convert_precision(args, hf_model, mg_model, template, test_convert_dtyp for n, m in mg_language_model.named_modules(): if n.endswith('router'): m.to(mg_dtype) - attention_flash_context = ( - _patch_attention_flash_fp32(mg_dtype) if args.attention_backend.name == 'flash' else nullcontext()) + attention_context = ( + _patch_attention_fp32(mg_dtype) if args.attention_backend.name in {'flash', 'fused'} else nullcontext()) with torch.inference_mode(), _model_cpu_forward_context( mg_modules, test_convert_dtype, 'cuda', share_embedding=share_embedding, - target_device=mg_device), attention_flash_context: + target_device=mg_device), attention_context: mg_logits = forward_step_helper(mg_model, mg_inputs, dtype=test_convert_dtype) if args.tensor_model_parallel_size > 1 and args.task_type != 'seq_cls': if mg_logits is not None: