Replacing --erf-gelu option with explicit --onnx-safe option

Signed-off-by: Boris Fomitchev <bfomitchev@nvidia.com>

megatron/arguments.py

@@ -158,9 +158,8 @@ def _add_network_size_args(parser):
                        help='Use OpenAIs GeLU implementation. This option'
                        'should not be used unless for backward compatibility'
                        'reasons.')
-    group.add_argument('--erf-gelu', action='store_true',
-                       help='Python GeLU implementation equivalent to one in Torch. This option'
-                       'should only be used to work around Torch bug exporting gelu() to ONNX in FP16')
+    group.add_argument('--onnx-safe', action='store_true',
+                       help='Use workarounds for known problems with Torch ONNX exporter')
 
     return parser
 

megatron/model/bert_model.py

@@ -95,8 +95,7 @@ class BertLMHead(MegatronModule):
         self.gelu = torch.nn.functional.gelu
         if args.openai_gelu:
             self.gelu = openai_gelu
-        # make it override 
-        if args.erf_gelu:
+        elif args.onnx_safe:
             self.gelu = erf_gelu
 
     def forward(self, hidden_states, word_embeddings_weight):

megatron/model/language_model.py

@@ -52,7 +52,7 @@ def get_language_model(attention_mask_func, num_tokentypes, add_pooler,
     gelu = F.gelu
     if args.openai_gelu:
         gelu = openai_gelu
-    if args.erf_gelu:
+    elif args.onnx_safe:
         gelu = erf_gelu
     
     # Language model.