From 763469830cf953d027f24af600f446fe26dcdc41 Mon Sep 17 00:00:00 2001 From: Qing Lan Date: Fri, 10 Feb 2023 01:30:57 -0600 Subject: [PATCH] place multi-processing init to main method (#443) --- examples/pytorch/gpt/utils/huggingface_gpt_convert.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/pytorch/gpt/utils/huggingface_gpt_convert.py b/examples/pytorch/gpt/utils/huggingface_gpt_convert.py index 3eb83006e..02c1a621b 100644 --- a/examples/pytorch/gpt/utils/huggingface_gpt_convert.py +++ b/examples/pytorch/gpt/utils/huggingface_gpt_convert.py @@ -158,9 +158,7 @@ def split_and_convert(args): "mlp.dense_4h_to_h.bias", "mlp.dense_4h_to_h.weight", ] - - torch.multiprocessing.set_start_method("spawn") - torch.multiprocessing.set_sharing_strategy("file_system") + pool = multiprocessing.Pool(args.processes) for name, param in model.named_parameters(): if name.find("weight") == -1 and name.find("bias") == -1: @@ -201,4 +199,6 @@ def split_and_convert(args): print("{}: {}".format(key, vars(args)[key])) print("========================================") + torch.multiprocessing.set_start_method("spawn") + torch.multiprocessing.set_sharing_strategy("file_system") split_and_convert(args)