Move flash_attn assert from __init__ into calling func (#32)

- Change the assert to warning in __init__ (e4f4e18f099e6fd012bfed0bd3170dc2ea209592)
- Update modeling_phi3_small.py (e203eff1c173b6ff6f356c6a9e9e9fb9f2371c72)
- Update modeling_phi3_small.py (bef256d02055a0870d55b9d21fc35d68bdee1c61)
- Moved flash_attn assert close to the caller (e4770004838da0fec2d9f8cdb60f8417e5394ba7)


Co-authored-by: Roger Feng <[email protected]>

modeling_phi3_small.py

@@ -215,7 +215,6 @@ class Phi3SmallSelfAttention(nn.Module):
                 f"Layer {layer_idx + 1} is using dense attention since it is divisible by "
                 f"{self.config.dense_attention_every_n_layers}"
             )
-            assert is_flash_attention_available, "Flash Attention is not available, but is needed for dense attention"
         else:
             # BlockSparse related Parameters
             self.blocksparse_params = BlockSparseParams.from_config(config)
@@ -419,6 +418,8 @@ class Phi3SmallSelfAttention(nn.Module):
             avoid doing that.
 
         """
+        assert is_flash_attention_available, "Flash Attention is not available, but is needed for dense attention"
+        
         attention_dropout_prob = self.attention_dropout_rate if self.training else 0.0
         # Get into the correct shape for the Flash Attention API
         # shape: (bs, seq_len, nqp, hn)