Update neox_args.py

These attention configuration options were missing from the docs. This will fix that.
EleutherAI · Nov 16, 2023 · b18f25c · b18f25c
1 parent d8028f8
commit b18f25c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/megatron/neox_arguments/neox_args.py b/megatron/neox_arguments/neox_args.py
@@ -175,7 +175,7 @@ class NeoXArgsModel(NeoXArgsTemplate):
     The first item in the list specifies the attention type(s), and should be a list of strings. The second item
     specifies the number of times to repeat those attention types in the full list.
 
-    attention type choices:  [global, local, sparse_fixed, sparse_variable, bslongformer, bigbird]
+    attention type choices:  [global, local, sparse_fixed, sparse_variable, bslongformer, bigbird, "gmlp", "amlp", "flash"]
 
     So a 12 layer network with only global attention could be specified like:
         [[[`global`], 12]]