Updating HPO for Perlmutter (ORNL#239)

* perlmutter

* perlmutter

* add perlmutter

* update on perlmutter

* add perlmutter hpo

* black

* add negative loss for deephyper

---------

Co-authored-by: Jong Choi <[email protected]>
Co-authored-by: Jong Choi <[email protected]>
Co-authored-by: Jong Choi <[email protected]>

Author: 4 people

examples/multidataset_hpo/gfm.py

@@ -60,6 +60,12 @@ def main():
     parser.add_argument(
         "--multi_model_list", help="multidataset list", default="OC2020"
     )
+    parser.add_argument(
+        "--num_samples",
+        type=int,
+        help="set num samples per process for weak-scaling test",
+        default=None,
+    )
 
     group = parser.add_mutually_exclusive_group()
     group.add_argument(
@@ -304,6 +310,13 @@ def main():
         ## Set local set
         for dataset in [trainset, valset, testset]:
             rx = list(nsplit(range(len(dataset)), local_comm_size))[local_comm_rank]
+            if args.num_samples is not None:
+                if args.num_samples > len(rx):
+                    log(
+                        f"WARN: requested samples are larger than what is available. Use only {len(rx)}."
+                    )
+                rx = rx[: args.num_samples]
+
             dataset.setkeys(common_variable_names)
             dataset.setsubset(rx[0], rx[-1] + 1, preload=True)
 

examples/multidataset_hpo/gfm_deephyper_multi_perlmutter.py

@@ -0,0 +1,177 @@
+import os, sys
+
+import torch
+
+torch.backends.cudnn.enabled = False
+
+# deprecated in torch_geometric 2.0
+try:
+    from torch_geometric.loader import DataLoader
+except:
+    from torch_geometric.data import DataLoader
+
+import pandas as pd
+import subprocess
+import re
+
+pd.options.display.max_columns = None
+pd.options.display.max_rows = None
+pd.options.display.width = None
+
+# Retrieve constants
+NNODES = int(os.environ["NNODES"])
+NTOTGPUS = int(os.environ["NTOTGPUS"])
+NNODES_PER_TRIAL = int(os.environ["NNODES_PER_TRIAL"])
+NGPUS_PER_TRIAL = int(os.environ["NGPUS_PER_TRIAL"])
+NUM_CONCURRENT_TRIALS = int(os.environ["NUM_CONCURRENT_TRIALS"])
+NTOT_DEEPHYPER_RANKS = int(os.environ["NTOT_DEEPHYPER_RANKS"])
+OMP_NUM_THREADS = int(os.environ["OMP_NUM_THREADS"])
+DEEPHYPER_LOG_DIR = os.environ["DEEPHYPER_LOG_DIR"]
+DEEPHYPER_DB_HOST = os.environ["DEEPHYPER_DB_HOST"]
+SLURM_JOB_ID = os.environ["SLURM_JOB_ID"]
+
+
+def _parse_results(stdout):
+    pattern = r"Val Loss: ([-+]?(\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?)"
+    matches = re.findall(pattern, stdout.decode())
+    if matches:
+        return matches[-1][0]
+    else:
+        return "F"
+
+
+def run(trial, dequed=None):
+    f = open(f"output-{SLURM_JOB_ID}-{trial.id}.txt", "w")
+    python_exe = sys.executable
+    python_script = os.path.join(os.path.dirname(__file__), "gfm.py")
+
+    # TODO: Launch a subprocess with `srun` to train neural networks
+    params = trial.parameters
+    log_name = "gfm" + "_" + str(trial.id)
+    master_addr = f"HYDRAGNN_MASTER_ADDR={dequed[0]}"
+    nodelist = ",".join(dequed)
+
+    # time srun -u -n32 -c2 --ntasks-per-node=8 --gpus-per-node=8 --gpu-bind=closest
+    prefix = " ".join(
+        [
+            f"srun",
+            f"-N {NNODES_PER_TRIAL} -n {NGPUS_PER_TRIAL}",
+            f"--ntasks-per-node=4 --gpus-per-node=4",
+            f"--cpus-per-task {OMP_NUM_THREADS} --threads-per-core 1 --cpu-bind threads",
+            f"--gpus-per-task=1",
+            f"--export=ALL,{master_addr},HYDRAGNN_MAX_NUM_BATCH=100,HYDRAGNN_USE_VARIABLE_GRAPH_SIZE=1,HYDRAGNN_AGGR_BACKEND=mpi",
+            f"--nodelist={nodelist}",
+            f"--output {DEEPHYPER_LOG_DIR}/output_{SLURM_JOB_ID}_{trial.id}.txt",
+            f"--error {DEEPHYPER_LOG_DIR}/error_{SLURM_JOB_ID}_{trial.id}.txt",
+        ]
+    )
+
+    command = " ".join(
+        [
+            prefix,
+            python_exe,
+            "-u",
+            python_script,
+            f"--model_type={trial.parameters['model_type']}",
+            f"--hidden_dim={trial.parameters['hidden_dim']}",
+            f"--num_conv_layers={trial.parameters['num_conv_layers']}",
+            f"--num_headlayers={trial.parameters['num_headlayers']}",
+            f"--dim_headlayers={trial.parameters['dim_headlayers']}",
+            f"--multi",
+            f"--ddstore",
+            # f'--multi_model_list="ANI1x,MPTrj,OC2020-2M,OC2022,qm7x"',
+            ## debugging
+            f'--multi_model_list="ANI1x,MPTrj,qm7x"',
+            f"--num_samples=3200",
+            f"--num_epoch=5",
+            f"--log={log_name}",
+        ]
+    )
+    print("Command = ", command, flush=True, file=f)
+
+    result = subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT)
+    output = "F"
+    try:
+        pattern = r"Val Loss: ([-+]?(\d+(\.\d*)?|\.\d+)([eE][-+]?\d+)?)"
+        fout = open(f"{DEEPHYPER_LOG_DIR}/error_{SLURM_JOB_ID}_{trial.id}.txt", "r")
+        while True:
+            line = fout.readline()
+            matches = re.findall(pattern, line)
+            if matches:
+                output = -float(matches[-1][0])
+            if not line:
+                break
+        fout.close()
+
+    except Exception as excp:
+        print(excp, flush=True, file=f)
+        output = "F"
+
+    print("Output:", output, flush=True, file=f)
+    objective = output
+    print(objective, flush=True, file=f)
+    metadata = {"some_info": "some_value"}
+    f.close()
+
+    return {"objective": objective, "metadata": metadata}
+
+
+if __name__ == "__main__":
+
+    log_name = f"gfm-{SLURM_JOB_ID}"
+
+    # Choose the sampler (e.g., TPESampler or RandomSampler)
+    from deephyper.evaluator import Evaluator, ProcessPoolEvaluator, queued
+    from deephyper.problem import HpProblem
+    from deephyper.search.hps import CBO
+    from hydragnn.utils.deephyper import read_node_list
+
+    # define the variable you want to optimize
+    problem = HpProblem()
+
+    # Define the search space for hyperparameters
+    problem.add_hyperparameter((2, 6), "num_conv_layers")  # discrete parameter
+    problem.add_hyperparameter((100, 2000), "hidden_dim")  # discrete parameter
+    problem.add_hyperparameter((1, 3), "num_headlayers")  # discrete parameter
+    problem.add_hyperparameter((100, 1000), "dim_headlayers")  # discrete parameter
+    problem.add_hyperparameter(
+        ["EGNN", "SchNet", "PNA"], "model_type"
+    )  # categorical parameter
+
+    # Create the node queue
+    queue, _ = read_node_list()
+    print("The queue:", queue, len(queue))
+    print("NNODES_PER_TRIAL", NNODES_PER_TRIAL)
+    print("NUM_CONCURRENT_TRIALS", NUM_CONCURRENT_TRIALS)
+    print("NGPUS_PER_TRIAL", NGPUS_PER_TRIAL)
+    print("NTOTGPUS", NTOTGPUS)
+    print(NTOTGPUS, NGPUS_PER_TRIAL, NTOTGPUS // NGPUS_PER_TRIAL, len(queue))
+
+    # Define the search space for hyperparameters
+    # define the evaluator to distribute the computation
+    evaluator = queued(ProcessPoolEvaluator)(
+        run,
+        num_workers=NUM_CONCURRENT_TRIALS,
+        queue=queue,
+        queue_pop_per_task=NNODES_PER_TRIAL,  # Remove the hard-coded value later
+    )
+
+    # Define the search method and scalarization
+    # search = CBO(problem, parallel_evaluator, random_state=42, log_dir=log_name)
+    search = CBO(
+        problem,
+        evaluator,
+        acq_func="UCB",
+        multi_point_strategy="cl_min",  # Constant liar strategy
+        random_state=42,
+        # Location where to store the results
+        log_dir=log_name,
+        # Number of threads used to update surrogate model of BO
+        n_jobs=OMP_NUM_THREADS,
+    )
+
+    timeout = None
+    results = search.search(max_evals=100, timeout=timeout)
+    print(results)
+
+    sys.exit(0)

examples/multidataset_hpo/job-perlmutter.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+#SBATCH -A m4452_g
+#SBATCH -J HydraGNN
+#SBATCH -C gpu
+#SBATCH -q regular
+#SBATCH -t 0:30:00
+#SBATCH -N 64
+##SBATCH --ntasks-per-node=4
+##SBATCH --gpus-per-task=1
+##SBATCH -c 32
+
+set -x
+
+export MIOPEN_DISABLE_CACHE=1
+#export HSA_DISABLE_CACHE=1
+
+#export ROCM_HOME=/opt/rocm-5.4.2
+# export TRANSFORMERS_OFFLINE=1
+# export HF_DATASETS_OFFLINE=1
+# export NCCL_DEBUG=INFO
+# export settings
+#export TORCH_EXTENSIONS_DIR=$PWD/deepspeed
+# export HF_HOME=$PWD/hfdata
+
+# setup hostfile
+HOSTS=.hosts-job$SLURM_JOB_ID
+HOSTFILE=hostfile.txt
+srun hostname > $HOSTS
+sed 's/$/ slots=4/' $HOSTS > $HOSTFILE
+
+# setup env file
+#echo "PATH=$PATH" > .deepspeed_env
+#echo "LD_LIBRARY_PATH=$LD_LIBRARY_PATH" >> .deepspeed_env
+#echo "CPATH=$CPATH" >> .deepspeed_env
+#echo "TORCH_EXTENSIONS_DIR=$PWD/deepspeed" >> .deepspeed_env
+#echo "HF_HOME=$PWD/hfdata" >> .deepspeed_env
+#echo "ROCM_HOME=/opt/rocm-5.4.0" >> .deepspeed_env
+
+# Configuration 
+export NNODES=$SLURM_JOB_NUM_NODES # e.g., 100 total nodes
+export NNODES_PER_TRIAL=64
+export NUM_CONCURRENT_TRIALS=$(( $NNODES / $NNODES_PER_TRIAL ))
+export NTOTGPUS=$(( $NNODES * 4 )) # e.g., 800 total GPUs
+export NGPUS_PER_TRIAL=$(( 4 * $NNODES_PER_TRIAL )) # e.g., 32 GPUs per training
+export NTOT_DEEPHYPER_RANKS=$(( $NTOTGPUS / $NGPUS_PER_TRIAL )) # e.g., 25 total DH ranks
+export OMP_NUM_THREADS=4 # e.g., 8 threads per rank
+[ $NTOTGPUS -ne $(($NGPUS_PER_TRIAL*$NUM_CONCURRENT_TRIALS)) ] && echo "ERROR!!" 
+
+#export CUDA_DEVICE_MAX_CONNECTIONS=1
+#export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+
+# DeepHyper variables
+export DEEPHYPER_LOG_DIR="deephyper-experiment"-$SLURM_JOB_ID 
+mkdir -p $DEEPHYPER_LOG_DIR
+export DEEPHYPER_DB_HOST=$HOST
+# Start Redis server (shared memory between search processes)
+# TODO: install Redis and set the `redis.conf` path here
+#export REDIS_CONF=...
+#pushd $DEEPHYPER_LOG_DIR
+#redis-server $REDIS_CONF &
+#popd
+
+# Safe sleep to let everything start
+sleep 5
+
+echo "Doing something"
+
+# Launch DeepHyper (1 rank per node, NTOT_DEEPHYPER_RANKS <= NNODES here)
+# meaning NGPUS_PER_TRAINING >= 8
+#$NTOT_DEEPHYPER_RANKS 
+#srun -n1 python qm9_deephyper_multi.py
+export HYDRAGNN_SYSTEM=perlmutter
+python gfm_deephyper_multi_perlmutter.py

examples/open_catalyst_2020/train.py

@@ -138,6 +138,7 @@ def get(self, idx):
     parser.add_argument("--shmem", action="store_true", help="shmem")
     parser.add_argument("--log", help="log name")
     parser.add_argument("--batch_size", type=int, help="batch_size", default=None)
+    parser.add_argument("--num_epoch", type=int, help="num_epoch", default=None)
     parser.add_argument("--everyone", action="store_true", help="gptimer")
     parser.add_argument("--modelname", help="model name")
 
@@ -181,6 +182,9 @@ def get(self, idx):
     if args.batch_size is not None:
         config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size
 
+    if args.num_epoch is not None:
+        config["NeuralNetwork"]["Training"]["num_epoch"] = args.num_epoch
+
     ##################################################################################################################
     # Always initialize for multi-rank training.
     comm_size, rank = hydragnn.utils.setup_ddp()

hydragnn/train/train_validate_test.py

@@ -154,6 +154,9 @@ def train_validate_test(
             if epoch == 0:
                 tr.reset()
 
+        if int(os.getenv("HYDRAGNN_VALTEST", "1")) == 0:
+            continue
+
         val_loss, val_taskserr = validate(
             val_loader, model, verbosity, reduce_ranks=True
         )

hydragnn/utils/deephyper.py

@@ -13,17 +13,31 @@ def master_from_host(host):
 def read_node_list():
     node_list = os.environ["SLURM_NODELIST"]
     nodes = []
-    node_subsets = node_list[9:-1].split(",")
-    for subset in node_subsets:
-        if "-" in subset:
-            start, end = subset.split("-")
-            start, end = int(start), int(end)
-            for i in range(start, end + 1):
-                leading_zeros = "".join(["0"] * (5 - len(str(i))))
-                nodes.append(f"frontier{leading_zeros}{i}")
-        else:
-            nodes.append(f"frontier{subset}")
-    nodes_string = ",".join(nodes)
+    system = os.getenv("HYDRAGNN_SYSTEM", "frontier")
+    if system == "frontier":
+        node_subsets = node_list[9:-1].split(",")
+        for subset in node_subsets:
+            if "-" in subset:
+                start, end = subset.split("-")
+                start, end = int(start), int(end)
+                for i in range(start, end + 1):
+                    leading_zeros = "".join(["0"] * (5 - len(str(i))))
+                    nodes.append(f"frontier{leading_zeros}{i}")
+            else:
+                nodes.append(f"frontier{subset}")
+        nodes_string = ",".join(nodes)
+    elif system == "perlmutter":
+        node_subsets = node_list[4:-1].split(",")
+        for subset in node_subsets:
+            if "-" in subset:
+                start, end = subset.split("-")
+                start, end = int(start), int(end)
+                for i in range(start, end + 1):
+                    leading_zeros = "".join(["0"] * (6 - len(str(i))))
+                    nodes.append(f"nid{leading_zeros}{i}")
+            else:
+                nodes.append(f"nid{subset}")
+        nodes_string = ",".join(nodes)
     return nodes, nodes_string