From 46c4695421d2102b662124a53a06edca6b3363a3 Mon Sep 17 00:00:00 2001 From: Chen Weihang Date: Mon, 4 Jan 2021 19:59:43 -0600 Subject: [PATCH] Set FLAGS_selected_gpus for spawn (#29962) * set flags_selectedd_gpus for spawn * add cond for unittest * Delete test_no_single_process_using_multi_gpus_in_spawn.py * Update spawn.py * Update nccl_context.cc --- python/paddle/distributed/spawn.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/python/paddle/distributed/spawn.py b/python/paddle/distributed/spawn.py index 433662e8ebc..86ec18061c5 100644 --- a/python/paddle/distributed/spawn.py +++ b/python/paddle/distributed/spawn.py @@ -27,7 +27,7 @@ from paddle.device import get_device # deprecated module import from paddle.fluid import core -from paddle.fluid.framework import _cpu_num +from paddle.fluid.framework import _cpu_num, set_flags class ParallelEnvArgs(object): @@ -153,6 +153,12 @@ def _remove_risky_env(): def _set_trainer_env(env_dict): + # NOTE(chenweihang): [ Why need set FLAGS_selected_gpus here? ] + # When the child process starts, it will inherit the configuration of the + # main process and set the FLAGS once, but the environment variable has + # not been set at this time, which leads to the FLAGS_selected_gpus + # is keep same with mainprocess(usually empty), so manually update the flags here + set_flags({'FLAGS_selected_gpus': env_dict['FLAGS_selected_gpus']}) for var_name in env_dict: os.environ[var_name] = env_dict[var_name] -- GitLab