Skip to content

fix: set CUDA device before init_process_group; use_orig_params for FSDP v1 #65

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
sparticlesteve wants to merge 2 commits into
base: mlp_weight_branches
Choose a base branch
from
Draft

fix: set CUDA device before init_process_group; use_orig_params for FSDP v1 #65

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
21b285c
fix: set CUDA device before init_process_group in setup_ddp
sparticlesteve Apr 24, 2026
3b63250
feat: add --datadir CLI argument to gfm_mlip_all_mpnn.py
sparticlesteve Apr 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions examples/multidataset_hpo_sc26/gfm_mlip_all_mpnn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -177,6 +177,7 @@ def set_param_value(param, group="Architecture"):
parser.add_argument("--correlation", type=int, help="correlation", default=None)
parser.add_argument("--nvme", action="store_true", help="use NVME")
parser.add_argument("--startfrom", type=str, help="startfrom", default=None)
parser.add_argument("--datadir", type=str, help="path to directory containing .bp dataset files", default=None)

group = parser.add_mutually_exclusive_group()
group.add_argument(
Expand Down Expand Up @@ -209,7 +210,7 @@ def set_param_value(param, group="Architecture"):
node_feature_names = ["atomic_number", "cartesian_coordinates", "forces"]
node_feature_dims = [1, 3, 3]
dirpwd = os.path.dirname(os.path.abspath(__file__))
datadir = os.path.join(dirpwd, "dataset")
datadir = args.datadir if args.datadir is not None else os.path.join(dirpwd, "dataset")
##################################################################################################################
input_filename = os.path.join(dirpwd, args.inputfile)
##################################################################################################################
Expand Down Expand Up @@ -425,9 +426,7 @@ def set_param_value(param, group="Architecture"):
ndata_list = list()
pna_deg_list = list()
for model in modellist:
fname = os.path.join(
os.path.dirname(__file__), "./dataset/%s-v2.bp" % model
)
fname = os.path.join(datadir, "%s-v2.bp" % model)
with adios2_open(fname, "r", MPI.COMM_SELF) as f:
f.__next__()
ndata = f.read_attribute("trainset/ndata").item()
Expand Down Expand Up @@ -595,7 +594,7 @@ def set_param_value(param, group="Architecture"):
local_comm_rank = local_comm.Get_rank()
local_comm_size = local_comm.Get_size()

fname = os.path.join(os.path.dirname(__file__), "./dataset/%s-v2.bp" % mymodel)
fname = os.path.join(datadir, "%s-v2.bp" % mymodel)

## FIXME: only for Frontier NVME
bbpath = f"/mnt/bb/{os.getenv('USER')}"
Expand Down Expand Up @@ -823,9 +822,7 @@ def dojob(cmd):
## No DDStore. Each process opens multiple adios files.
filename_list = list()
for model in modellist:
fname = os.path.join(
os.path.dirname(__file__), "./dataset/%s-v2.bp" % model
)
fname = os.path.join(datadir, "%s-v2.bp" % model)
filename_list.append(fname)

kwargs = {"var_config": var_config, "keys": common_variable_names}
Expand Down
13 changes: 12 additions & 1 deletion hydragnn/utils/distributed/distributed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,16 @@ def _derive_master_port(default_port: int = 8889) -> str:
## Setting LOCAL_RANK complicts with DeviceMesh when using the srun "--gpus-per-task=1" option
# os.environ["LOCAL_RANK"] = str(get_local_rank())

# Set the CUDA device for this rank BEFORE init_process_group so
# that the NCCL communicator created by init_process_group (comm '0')
# uses the correct per-rank GPU. Without this, all ranks may use
# CUDA device 0, causing FSDP v2's per-rank device keys to mismatch
# comm '0' and trigger a new comm ('1') → store-key deadlock.
if backend == "nccl" and torch.cuda.is_available():
local_rank = get_local_rank()
if local_rank < torch.cuda.device_count():
torch.cuda.set_device(local_rank)

if (backend == "gloo") and ("GLOO_SOCKET_IFNAME" not in os.environ):
ifname = find_ifname(master_addr)
if ifname is not None:
Expand Down Expand Up @@ -393,6 +403,7 @@ def is_model_distributed(model):
return isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel)



def get_distributed_model(
model,
verbosity=0,
Expand Down Expand Up @@ -447,7 +458,7 @@ def get_distributed_model(
if use_fsdp:
if fsdp_version == 1:
print_distributed(verbosity, "Using FSDP v1 wrapper")
model = FSDP(model, sharding_strategy=sharding_strategy)
model = FSDP(model, sharding_strategy=sharding_strategy, use_orig_params=True)
else:
if fsdp_strategy not in ["FULL_SHARD", "SHARD_GRAD_OP"]:
raise ValueError(
Expand Down