try flash, no compile

allenai · Dec 3, 2024 · 5858777 · 5858777
1 parent 2863fe9
commit 5858777
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/scripts/train/OLMo2-26B.py b/src/scripts/train/OLMo2-26B.py
@@ -24,7 +24,9 @@
 def build_model_config(common: CommonComponents) -> TransformerConfig:
     return TransformerConfig.olmo2_26B(
         vocab_size=common.tokenizer.padded_vocab_size(),
-        compile=True,
+        compile=False,
+        fused_ops=False,
+        use_flash=True,
         dp_config=TransformerDataParallelConfig(
             name=DataParallelType.fsdp, param_dtype=DType.bfloat16, reduce_dtype=DType.float32
         ),