..

mosaicml · Dec 5, 2024 · 18afcc5 · 18afcc5
1 parent f1ad991
commit 18afcc5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/llmfoundry/models/layers/blocks.py b/llmfoundry/models/layers/blocks.py
@@ -51,6 +51,7 @@ def __init__(
     ):
         if attn_config is None:
             attn_config = attn_config_defaults
+            attn_config.pop('flex_attn_compile', None)
 
         if ffn_config is None:
             self.ffn_config: dict[str, Any] = {

diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -423,7 +423,7 @@ def __init__(self, config: MPTConfig):
         self.mb_args = None
         self.shift_labels = True
 
-        flex_attn_compile = config.attn_config.pop('flex_attn_compile')
+        flex_attn_compile = config.attn_config.pop('flex_attn_compile', False)
         if self.attn_impl == 'flex':
             self.compiled_flex_attention = torch.compile(
                 flex_attention,