implement multi-gpu training

we use the torchrun thingie, which limits us to single-node training

Author: pdobbelaere

psiflow/execution.py

@@ -246,7 +246,7 @@ def __init__(
             )
 
         if self.executor_type == "workqueue":
-            # WQ-specific checks
+            # WQ-specific checks  TODO: what about multinode?
             ensure(
                 self.kwargs["gpus_per_task"] <= resources["gpus"],
                 self.kwargs["cores_per_task"] <= resources["cores"],
@@ -541,15 +541,10 @@ def __init__(self, **kwargs) -> None:
                 "ModelTraining is configured for CPU operation. Is this what you want?"
             )
 
-        # if self.multigpu:
-        #     # TODO: why? Think this might be a multinode thing - which I do not care about
-        #     message = (
-        #         "the max_training_time keyword does not work "
-        #         "in combination with multi-gpu training. Adjust "
-        #         "the maximum number of epochs to control the "
-        #         "duration of training"
-        #     )
-        #     assert self.max_runtime is None, message
+    @property
+    def multi_gpu(self) -> bool:
+        # only for WQ
+        return (self.spec or {}).get("gpus", 0) > 1
 
     def train_command(self, initialize: bool = False):
         command = "psiflow-mace-train"

psiflow/models/mace.py

@@ -71,17 +71,6 @@ def format_E0s(atomic_energies: dict) -> dict | str:
         return "average"
 
 
-def get_bash_command(root: Path, pre: Optional[str] = None) -> str:
-    defaults = [
-        "mkdir checkpoints",  # otherwise MACE borks
-        "mace_run_train --config {}",  # run MACE
-        f"rsync -av --ignore-existing --exclude=/*.model ./ {root}/",  # copy things back
-    ]
-    if pre:
-        defaults.insert(1, pre)
-    return "\n".join(defaults)
-
-
 def _execute(
     bash_template: str,
     inputs: list[File],
@@ -176,19 +165,49 @@ def _resolve_futures(self) -> AppFuture:
         cfg = self.config | {KEY_ATOMIC_ENERGIES: self.atomic_energies}
         return resolve_nested_futures(cfg)
 
-    @staticmethod
-    def get_app(init: bool = False) -> tuple[Callable, str]:
-        # TODO: update -- why staticmethod and not function?
+    def _execute_app(self, config: dict) -> AppFuture:
+        """"""
         context = psiflow.context()
         definition = context.definitions["ModelTraining"]
-        template = context.bash_template
+        resources = definition.wq_resources()
+
+        # final config tweaks
+        if definition.multi_gpu:
+            config["distributed"] = True
+            config["launcher"] = "torchrun"
+        else:
+            config["distributed"] = False
+        file = psiflow.context().new_file("mace_cfg_", ".yaml")
+        yaml.safe_dump(config, open(file.filepath, "w"))
+
+        # construct MACE train script
+        command = "$(which mace_run_train) --config {}"
+        if config["distributed"]:
+            command = f"torchrun --standalone --nnodes=1 --nproc_per_node={resources['gpus']} {command}"
+        command = definition.wrap_in_timeout(command)
+
+        pre_copy = ""
+        if self.has_checkpoint:  # restart from latest checkpoint
+            chkpt_out = f"checkpoints/{create_checkpoint_name(config)}"
+            pre_copy = f"cp {self.path_chkpt} {chkpt_out}"
+
+        command_lines = [
+            "mkdir checkpoints",  # otherwise MACE borks
+            pre_copy,
+            command,
+            f"rsync -av --ignore-existing --exclude=/*.model ./ {self.root}/",  # copy things back
+        ]
+        command = "\n".join([l for l in command_lines if l])
+
+        execute_app = bash_app(_execute, executors=["ModelTraining"])
         env_vars = format_env_vars(definition.env_vars)
-        app = partial(
-            bash_app(_execute, executors=["ModelTraining"]),
-            parsl_resource_specification=definition.wq_resources(),
-            label="mace_init" if init else "mace_train",
+        future = execute_app(
+            bash_template=context.bash_template.format(commands=command, env=env_vars),
+            inputs=[file],
+            parsl_resource_specification=resources,
+            label="mace_init" if config["name"] == "init" else "mace_train",
         )
-        return app, template.format(commands="{commands}", env=env_vars)
+        return future
 
     @property
     def path_config(self) -> Path:
@@ -256,13 +275,8 @@ def initialize_app(
         "valid_file": file_val.filepath,
     }
     cfg["E0s"] = format_E0s(model.atomic_energies)
-    file_cfg = psiflow.context().new_file("mace_cfg_", ".yaml")
-    yaml.safe_dump(cfg, open(file_cfg.filepath, "w"))
-
-    command = get_bash_command(model.root)
-    app, template = model.get_app(init=True)
 
-    future = app(template.format(commands=command), inputs=[file_cfg])
+    future = model._execute_app(cfg)
     inputs = [future.stdout, future.stderr, future]
     future_ = process_output(model, config, inputs=inputs)
     return future_
@@ -292,17 +306,8 @@ def train_app(
         "valid_file": file_val.filepath,
     }
     cfg["E0s"] = format_E0s(cfg.pop(KEY_ATOMIC_ENERGIES))
-    file_cfg = psiflow.context().new_file("mace_cfg_", ".yaml")
-    yaml.safe_dump(cfg, open(file_cfg.filepath, "w"))
-
-    pre = None
-    if model.has_checkpoint:
-        chkpt_out = f"checkpoints/{create_checkpoint_name(cfg)}"
-        pre = f"cp {model.path_chkpt} {chkpt_out}"
-    command = get_bash_command(model.root, pre=pre)
-    app, template = model.get_app(init=False)
 
-    future = app(template.format(commands=command), inputs=[file_cfg])
+    future = model._execute_app(cfg)
     inputs = [future.stdout, future.stderr, future]
     future_ = process_output(model, config, inputs=inputs)
     return future_

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "psiflow"
-version = "4.0.3"
+version = "4.0.4"
 description = "Library for developing interatomic potentials"
 readme = "README.md"
 requires-python = ">=3.10"

tests/test_models.py

@@ -112,10 +112,13 @@ def test_mace_hamiltonian(dataset, mace_foundation):
     assert hamiltonian0 != hamiltonian1
     hamiltonian2 = psiflow.deserialize(psiflow.serialize(hamiltonian1)).result()
     assert hamiltonian0 != hamiltonian2
+    assert hamiltonian1 == hamiltonian2
+    hamiltonian2.update_kwargs(enable_cueq=True)
 
     e0 = hamiltonian0.compute(dataset, "energy")
     e1 = hamiltonian1.compute(dataset, "energy")
+    e2 = hamiltonian2.compute(dataset, "energy")
     assert np.allclose(e0.result(), e1.result())
-    pass
+    assert np.allclose(e0.result(), e2.result())