Update to work with args 2.

megatron/arguments.py

@@ -272,6 +272,14 @@ def _add_training_args(parser):
  help='Disables weight tying between embedding weights and final Linear layer') 
  group.add_argument('--sinusoidal-pos-emb', action='store_true',
  help='Uses Sinusoidal Positional embedding applied to the inputs instead of learned') 
+ group.add_argument('--rpe', action='store_true',
+ help='T5 relative positional encoding')
+ group.add_argument('--rpe-causal', action='store_true',
+ help='T5 relative positional encoding causal flag')
+ group.add_argument('--rpe-num-buckets', type=int, default=32,
+ help='T5 relative positional encoding number of buckets, default 32.')
+ group.add_argument('--rpe-max-distance', type=int, default=128,
+ help='T5 relative positional encoding max distance, default 128.')
  group.add_argument('--bias-dropout-fusion', action='store_true',
  help='Enable bias and dropout fusion.')
  group.add_argument('--sparsity', type=str, default='none',

megatron/model/transformer.py

@@ -591,10 +591,10 @@ def __init__(self, attention_mask_func,
  super(ParallelTransformer, self).__init__()
  args = get_args()
 
- self.rpe = args.rpe # True
- self.rpe_causal = args.rpe_causal # False
- self.rpe_num_buckets =.args.rpe_num_buckets # 32
- self.rpe_max_distance = args.rpe_max_distance # 128
+ rpe = args.rpe
+ rpe_causal = args.rpe_causal
+ rpe_num_buckets = args.rpe_num_buckets
+ rpe_max_distance = args.rpe_max_distance
 
  # Store activation checkpoiting flag.
  self.checkpoint_activations = args.checkpoint_activations