removed flash attention

models/micro_llama.py

@@ -249,7 +249,7 @@ class MiCRoLlama(LlamaPreTrainedModel, GenerationMixin):
         self.config: MiCRoLlamaConfig = config
         self.config.torch_dtype = torch.bfloat16
         self.config.use_bfloat16 = True
-        self.config._attn_implementation = "flash_attention_2" # {sdpa, flash_attention_2, eager}
+        self.config._attn_implementation = "eager" # {sdpa, flash_attention_2, eager}
         self.config.backbone_num_layers = self.config.num_hidden_layers
         self.config.num_hidden_layers = self.config.num_hidden_layers * run_config["num-experts"]
         self.config.loss_type = "ForCausalLMLoss"

models/micro_moe_llama.py

@@ -275,7 +275,7 @@ class MiCRoLlamaMoE(LlamaPreTrainedModel, GenerationMixin):
         self.config: MiCRoLlamaMoEConfig = config
         self.config.torch_dtype = torch.bfloat16
         self.config.use_bfloat16 = True
-        self.config._attn_implementation = "flash_attention_2" # {sdpa, flash_attention_2, eager}
+        self.config._attn_implementation = "eager" # {sdpa, flash_attention_2, eager}
         self.config.use_cache = True
         self.config.backbone_num_layers = self.config.num_hidden_layers
         self.config.num_hidden_layers = self.config.num_hidden_layers

models/micro_olmo.py

@@ -191,7 +191,7 @@ class MiCRoOLMo(Olmo2PreTrainedModel, GenerationMixin):
         self.config: Olmo2Config = config
         self.config.torch_dtype = torch.bfloat16
         self.config.use_bfloat16 = True
-        self.config._attn_implementation = "flash_attention_2" # {sdpa, flash_attention_2, eager}
+        self.config._attn_implementation = "eager" # {sdpa, flash_attention_2, eager}
         self.config.use_cache = True
         self.config.backbone_num_layers = self.config.num_hidden_layers
         self.config.num_hidden_layers = self.config.num_hidden_layers * run_config["num-experts"]

requirements.txt

@@ -3,5 +3,4 @@ plotly>=5.22.0
 pandas>=2.2.0
 torch==2.7.1
 transformers==4.53.2
-numpy==2.3.4
-flash-attn
+numpy==2.3.4

router_backend.py

@@ -195,7 +195,7 @@ def build_model(model_id: str, hf_token: str, use_cache: bool = True):
 
     model_config.torch_dtype = torch.bfloat16
     model_config.use_bfloat16 = True
-    model_config._attn_implementation = "flash_attention_2"
+    model_config._attn_implementation = "eager" # {sdpa, flash_attention_2, eager}
     model_config.use_cache = use_cache
     model_config.ablate = []