Change some arguments to default to on.

megatron/arguments.py

@@ -183,13 +183,15 @@ def parse_args(extra_args_provider=None, defaults={},
  'for distribute-checkpointed-activations to work you '\
  'need to enable checkpoint-activations'
 
- # load scaled_upper_triang_masked_softmax_fusion kernel
- if args.scaled_upper_triang_masked_softmax_fusion:
- fused_kernels.load_scaled_upper_triang_masked_softmax_fusion_kernel()
-
- # load scaled_masked_softmax_fusion kernel
  if args.scaled_masked_softmax_fusion:
- fused_kernels.load_scaled_masked_softmax_fusion_kernel()
+ if args.scaled_upper_triang_masked_softmax_fusion:
+ fused_kernels.load_scaled_upper_triang_masked_softmax_fusion_kernel()
+ else:
+ fused_kernels.load_scaled_masked_softmax_fusion_kernel()
+ else:
+ # This argument will eventually go away, for now make sure it is off
+ # if scaled_masked_softmax_fusion is off.
+ args.scaled_upper_triang_masked_softmax_fusion = False
 
  # Load mixed precision fused layer norm.
  if args.fp32_residual_connection:
@@ -328,18 +330,22 @@ def _add_training_args(parser):
  help='Exit the program after this many minutes.')
  group.add_argument('--tensorboard-dir', type=str, default=None,
  help='Write TensorBoard logs to this directory.')
+ group.add_argument('--no-scaled-masked-softmax-fusion',
+ action='store_false',
+ help='Disable fusion of query_key_value scaling, '
+ 'masking, and softmax.',
+ dest='scaled_masked_softmax_fusion')
  group.add_argument('--scaled-upper-triang-masked-softmax-fusion',
- action='store_true',
- help='Enable fusion of query_key_value_scaling '
- 'time (upper diagonal) masking and softmax.')
- group.add_argument('--scaled-masked-softmax-fusion',
- action='store_true',
- help='Enable fusion of query_key_value_scaling '
- 'general masking and softmax.')
- group.add_argument('--bias-gelu-fusion', action='store_true',
- help='Enable bias and gelu fusion.')
- group.add_argument('--bias-dropout-fusion', action='store_true',
- help='Enable bias and dropout fusion.')
+ type=bool,
+ help='Use upper triangular version of fused '
+ 'scale, mask, softmax fusion kernel (default for GPT). '
+ '- DEPRECATED')
+ group.add_argument('--no-bias-gelu-fusion', action='store_false',
+ help='Disable bias and gelu fusion.',
+ dest='bias_gelu_fusion')
+ group.add_argument('--no-bias-dropout-fusion', action='store_false',
+ help='Disable bias and dropout fusion.',
+ dest='bias_dropout_fusion')
 
  return parser
 
@@ -447,12 +453,13 @@ def _add_mixed_precision_args(parser):
  help='hysteresis for dynamic loss scaling')
  group.add_argument('--fp32-residual-connection', action='store_true',
  help='Move residual connections to fp32.')
- group.add_argument('--apply-query-key-layer-scaling', action='store_true',
- help='Scale Q * K^T by 1 / layer-number. If this flag '
- 'is set, then it will automatically set '
- 'attention-softmax-in-fp32 to true')
+ group.add_argument('--no-query-key-layer-scaling', action='store_false',
+ help='Do not scale Q * K^T by 1 / layer-number.',
+ dest='apply_query_key_layer_scaling')
  group.add_argument('--attention-softmax-in-fp32', action='store_true',
- help='Run attention masking and softmax in fp32.')
+ help='Run attention masking and softmax in fp32. '
+ 'This flag is ignored unless '
+ '--no-query-key-layer-scaling is specified.')
  group.add_argument('--fp32-allreduce', action='store_true',
  help='All-reduce in fp32')
  group.add_argument('--fp16-lm-cross-entropy', action='store_true',

pretrain_gpt2.py

@@ -141,4 +141,5 @@ def train_valid_test_datasets_provider(train_val_test_num_samples):
 if __name__ == "__main__":
 
  pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
- args_defaults={'tokenizer_type': 'GPT2BPETokenizer'})
+ args_defaults={'tokenizer_type': 'GPT2BPETokenizer',
+ 'scaled_upper_triang_masked_softmax_fusion': True})