Change Num Partitions size expansion fix (NVIDIA#4719)

* add cloning Signed-off-by: Abhinav Khattar <[email protected]> * map to cpu Signed-off-by: Abhinav Khattar <[email protected]> Signed-off-by: Abhinav Khattar <[email protected]>
Zenodia · Aug 10, 2022 · 66c7677 · 66c7677
1 parent f921ebe
commit 66c7677
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/nlp/language_modeling/megatron_change_num_partitions.py b/examples/nlp/language_modeling/megatron_change_num_partitions.py
@@ -95,7 +95,7 @@ def split_partition(model, partitions, tp_size, write_path=None):
 
         idx = 0
         for name, param in model.named_parameters():
-            split_val = splits[idx][i]
+            split_val = splits[idx][i].clone()
 
             if param.shape != split_val.shape:
                 logging.info(
@@ -178,7 +178,7 @@ def main():
             merge_partition(model, partitions, args.target_file)
     else:
         app_state.model_parallel_size = 1
-        model = cls.restore_from(restore_path=args.model_file, trainer=trainer)
+        model = cls.restore_from(restore_path=args.model_file, trainer=trainer, map_location=torch.device("cpu"))
 
     if tgt_tp_size > 1:
         partitions = []