fix?

mosaicml · Jan 10, 2024 · 6550af9 · 6550af9
1 parent 350f6ca
commit 6550af9
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/callbacks/test_hf_checkpointer.py b/tests/callbacks/test_hf_checkpointer.py
@@ -30,7 +30,10 @@ def _dummy_log_info(*msgs: str):
     lambda *_, **__: None)
 def assert_checkpoint_saves_to_uri(uri: str, build_tiny_hf_mpt: Callable):
     uri_base = uri.split('://')[0]
-    model = build_tiny_hf_mpt()
+    model = build_tiny_hf_mpt(attn_config={
+        'attn_impl': 'triton',
+        'attn_uses_sequence_id': False,
+    })
 
     dummy_state = State(model=model,
                         rank_zero_seed=42,