Upload modeling_cloverlm.py with huggingface_hub

modeling_cloverlm.py

@@ -111,15 +111,27 @@ class MHSA(nn.Module):
 
             dtype = Q.dtype if Q.dtype in (torch.bfloat16, torch.float16) else torch.bfloat16
             if attn_backend == "flash2":
-                import flash_attn
+                try:
+                    import flash_attn
+                except ImportError as e:
+                    e.add_note(f"Can't run `attn_backend=flash2` because can't import flash_attn")
+                    raise e
                 Y = flash_attn.flash_attn_func(Q.to(dtype), K.to(dtype), V.to(dtype), causal=True, softmax_scale=1.0)
             elif attn_backend == "flash3":
                 import importlib
-                _fa3 = importlib.import_module("flash_attn_interface")
+                try:
+                    _fa3 = importlib.import_module("flash_attn_interface")
+                except ImportError as e:
+                    e.add_note(f"Can't run `attn_backend=flash3` because can't import flash_attn_interface")
+                    raise e
                 Y = _fa3.flash_attn_func(Q.to(dtype), K.to(dtype), V.to(dtype), causal=True, softmax_scale=1.0)
             elif attn_backend == "flash4":
                 import importlib
-                _fa4 = importlib.import_module("flash_attn.cute")
+                try:
+                    _fa4 = importlib.import_module("flash_attn.cute")
+                except ImportError as e:
+                    e.add_note(f"Can't run `attn_backend=flash4` because can't import flash_attn.cute")
+                    raise e
                 Y = _fa4.flash_attn_func(Q.to(dtype), K.to(dtype), V.to(dtype), causal=True, softmax_scale=1.0)[0]
             Y = Y.to(Q.dtype).flatten(-2, -1)