fix ulysses with vl (#5391)

Author: tastelikefeet

swift/trainers/sequence_parallel/ulysses.py

@@ -240,7 +240,7 @@ def pre_forward_split_hook(_self, args, kwargs):
 
         def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
                              dist_attn, **kwargs):
-            if module not in text_model.modules():
+            if module.__class__ not in [m.__class__ for m in text_model.modules()]:
                 return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states,
                                                                            value_states, attention_mask, *args,
                                                                            **kwargs)
@@ -261,7 +261,7 @@ def _attention(query, key, value, *args, **kwargs):
 
         def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
                             dist_attn, **kwargs):
-            if module not in text_model.modules():
+            if module.__class__ not in [m.__class__ for m in text_model.modules()]:
                 return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states,
                                                               attention_mask, *args, **kwargs)
             if dist_attn.local_attn is None: