Change the default setup

microsoft · Oct 13, 2023 · cd4e4c6 · cd4e4c6
1 parent 4c4bc76
commit cd4e4c6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/llama/model_test.py b/examples/llama/model_test.py
@@ -429,15 +429,15 @@ def test(args, batch_size, seq_len, dtype, world_size):
     # Configurations
     args = ModelArgs7B()
     batch_size = 1
-    seq_len = 2048
+    seq_len = 1024
     dtype = np.float16
     world_size = 1
 
     # Default from HuggingFace
     args.vocab_size = 32000
 
     # PyTorch model cannot run all layers due to OOM
-    args.n_layers = 24
+    args.max_seq_len = 1024
 
     # Verify the configurations
     assert batch_size <= args.max_batch_size