fix: gracefully fall back flash_attention_2 -> sdpa when flash_attn not installed

gpt2_decoder.py

@@ -99,6 +99,8 @@ class MossTTSNanoGPT2Attention(nn.Module):
         self.head_dim = hidden_size // num_heads
         self.embed_dim = hidden_size
         self.layer_idx = layer_idx
+        if attn_implementation == "flash_attention_2" and not _FLASH_ATTN_AVAILABLE:
+            attn_implementation = "sdpa"  # graceful fallback when flash_attn isn't installed
         self.attn_implementation = attn_implementation
         self.attn_dropout = float(config.attn_pdrop)
         self.resid_dropout = nn.Dropout(config.resid_pdrop)
@@ -388,6 +390,8 @@ class MossTTSNanoGPT2Model(nn.Module):
     def __init__(self, config: GPT2Config, attn_implementation: str = "eager") -> None:
         super().__init__()
         self.config = config
+        if attn_implementation == "flash_attention_2" and not _FLASH_ATTN_AVAILABLE:
+            attn_implementation = "sdpa"  # graceful fallback when flash_attn isn't installed
         self.attn_implementation = attn_implementation
         self.position_embedding_type = str(getattr(config, "position_embedding_type", "absolute")).lower()
         if self.position_embedding_type not in {"absolute", "rope"}: