fix gpt3 e2e test

google · Feb 15, 2024 · 2261012 · 2261012
1 parent 4c9eac6
commit 2261012
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/MaxText/configs/models/gpt3-52k.yml b/MaxText/configs/models/gpt3-52k.yml
@@ -37,3 +37,4 @@ adam_b1: 0.9
 adam_b2: 0.95
 adam_eps: 1.e-8
 adam_weight_decay: 0.1
+attention: "dot_product" # head_dim 8 is too small for splash/flash attention
diff --git a/MaxText/tests/gpt3_test.py b/MaxText/tests/gpt3_test.py
@@ -60,7 +60,6 @@ def setUp(self):
  super().setUp()
  pyconfig.initialize(
  [sys.argv[0], 'configs/base.yml'],
- attention="dot_product",
  run_name='test',
  enable_checkpointing=False,
  model_name='gpt3-52k',