diff --git a/examples/multidataset_hpo_sc26/gfm_mlip_all_mpnn.py b/examples/multidataset_hpo_sc26/gfm_mlip_all_mpnn.py
index 38872d762..812bd8c74 100644
--- a/examples/multidataset_hpo_sc26/gfm_mlip_all_mpnn.py
+++ b/examples/multidataset_hpo_sc26/gfm_mlip_all_mpnn.py
@@ -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(
@@ -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)
     ##################################################################################################################
@@ -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()
@@ -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')}"
@@ -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}
diff --git a/hydragnn/utils/distributed/distributed.py b/hydragnn/utils/distributed/distributed.py
index 124cdbbe5..201e8b7c8 100644
--- a/hydragnn/utils/distributed/distributed.py
+++ b/hydragnn/utils/distributed/distributed.py
@@ -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:
@@ -393,6 +403,7 @@ def is_model_distributed(model):
     return isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel)
 
 
+
 def get_distributed_model(
     model,
     verbosity=0,
@@ -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(