Address comments, fix argument bug.

megatron/arguments.py

@@ -503,7 +503,7 @@ def _add_distributed_args(parser):
                        ' and returns function to complete it instead.'
                        'Also turns on --use-cpu-initialization flag.'
                        'This is for external DDP manager.' )
-    group.add_argument('--use-cpu-initialization', type=bool, required=False,
+    group.add_argument('--use-cpu-initialization', action='store_true', default=None,
                        help='If set, affine parallel weights initialization uses CPU' )
     return parser
 

megatron/mpu/layers.py

@@ -260,9 +260,7 @@ class ColumnParallelLinear(torch.nn.Module):
                     self.output_size_per_partition,
                     device=torch.cuda.current_device(),
                     dtype=args.params_dtype))
-            self.bias.tensor_model_parallel = True
-            self.bias.partition_dim = 0
-            self.bias.partition_stride = stride
+            set_tensor_model_parallel_attributes(self.bias, True, 0, stride)
             # Always initialize bias to zero.
             with torch.no_grad():
                 self.bias.zero_()

tools/merge_mp_partitions.py

@@ -199,15 +199,16 @@ def main():
                                           'no_load_rng': True,
                                           'save_interval': 1})
     args = get_args()
-    model_type = args.model_type
-    orig_tensor_model_parallel_size = args.tensor_model_parallel_size
-    args.tensor_model_parallel_size = 1
-    tokenizer = rebuild_tokenizer(args)
 
     if args.pipeline_model_parallel_size > 1:
         print("Checkpoints with pipeline model parallelism are not currently supported.")
         exit()
 
+    model_type = args.model_type
+    orig_tensor_model_parallel_size = args.tensor_model_parallel_size
+    args.tensor_model_parallel_size = 1
+    tokenizer = rebuild_tokenizer(args)
+
     print('\n merging model parallel partitions ...')
     print(' > number of partitions: {}'.format(orig_tensor_model_parallel_size))
     print(' > checkpoint path: {}'.format(args.load))