Dev/multiple gpu

qmp/algorithms/chop_imag.py

@@ -6,11 +6,20 @@
 import typing
 import dataclasses
 import omegaconf
+import torch
 import torch.utils.tensorboard
 from ..utility.context import RuntimeContext
 from ..utility.action_dict import action_dict
 
 
+def _get_devices_from_context(context: RuntimeContext) -> list[torch.device]:
+    """
+    Get the devices list from the runtime context.
+    """
+    assert context.devices is not None
+    return context.devices
+
+
 @dataclasses.dataclass
 class ChopImagConfig:
     """
@@ -35,6 +44,9 @@ def main(
         model = context.create_model(runtime_config.model)
         data = checkpoint_data
 
+        # Get devices for multi-GPU computation
+        devices = _get_devices_from_context(context)
+
         logging.info(
             "Arguments Summary: Chop Size: %d, Second Order Magnitude: %.10f",
             self.chop_size,
@@ -59,7 +71,7 @@ def main(
             logging.info("The number of configurations: %d", num_configs)
             writer.add_scalar("chop_imag/num_configs", num_configs, i)  # type: ignore[no-untyped-call]
             psi = psi / psi.norm()
-            hamiltonian_psi = model.apply_within(configs, psi, configs)
+            hamiltonian_psi = model.apply_within(configs, psi, configs, devices)
             psi_hamiltonian_psi = (psi.conj() @ hamiltonian_psi).real
             energy = psi_hamiltonian_psi
             logging.info(

qmp/algorithms/guide.py

@@ -12,6 +12,14 @@
 from ..utility.action_dict import action_dict
 
 
+def _get_devices_from_context(context: RuntimeContext) -> list[torch.device]:
+    """
+    Get the devices list from the runtime context.
+    """
+    assert context.devices is not None
+    return context.devices
+
+
 @dataclasses.dataclass
 class GuideConfig:
     """
@@ -50,6 +58,9 @@ def main(
 
         data = checkpoint_data
 
+        # Get devices for multi-GPU computation
+        devices = _get_devices_from_context(context)
+
         # Create Optimizers
         # We use the same optimizer configuration for both network and sampling
         optimizer_network = context.create_optimizer(
@@ -96,7 +107,7 @@ def main(
                     [
                         configs_src_network,
                         model.find_relative(
-                            configs_src_network, psi_src_network, self.relative_count - len(configs_src_network)
+                            configs_src_network, psi_src_network, self.relative_count - len(configs_src_network), devices=devices
                         ),
                     ]
                 )
@@ -107,7 +118,7 @@ def main(
                     [
                         configs_src_sampling,
                         model.find_relative(
-                            configs_src_sampling, psi_src_sampling, self.relative_count - len(configs_src_sampling)
+                            configs_src_sampling, psi_src_sampling, self.relative_count - len(configs_src_sampling), devices=devices
                         ),
                     ]
                 )
@@ -118,7 +129,7 @@ def energy_sampling() -> torch.Tensor:
                 psi_src = network(configs_src)
                 with torch.no_grad():
                     psi_dst = network(configs_dst)
-                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src)
+                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src, devices)
                 psi_src_reweight = psi_src.conj() * reweight_sampling
                 num = psi_src_reweight @ hamiltonian_psi_dst
                 den = psi_src_reweight @ psi_src.detach()
@@ -132,7 +143,7 @@ def energy_network() -> torch.Tensor:
                     configs_dst = configs_dst_network
                     psi_src = network(configs_src)
                     psi_dst = network(configs_dst)
-                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src)
+                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src, devices)
                     psi_src_reweight = psi_src.conj() * reweight_network
                     num = psi_src_reweight @ hamiltonian_psi_dst
                     den = psi_src_reweight @ psi_src.detach()
@@ -145,7 +156,7 @@ def distribution() -> torch.Tensor:
                 with torch.no_grad():
                     psi_src = network(configs_src)
                     psi_dst = network(configs_dst)
-                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src)
+                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src, devices)
                     psi_src_reweight = psi_src.conj() * reweight_sampling
                     num = psi_src_reweight @ hamiltonian_psi_dst
                     den = psi_src_reweight @ psi_src

qmp/algorithms/haar.py

@@ -18,6 +18,14 @@
 from ..utility.optimizer import scale_learning_rate
 
 
+def _get_devices_from_context(context: RuntimeContext) -> list[torch.device]:
+    """
+    Get the devices list from the runtime context.
+    """
+    assert context.devices is not None
+    return context.devices
+
+
 @dataclasses.dataclass
 class _DynamicLanczos:
     """
@@ -35,6 +43,7 @@ class _DynamicLanczos:
     first_extend: bool
     eigen_count: int = 1
     period: int = 256
+    devices: list[torch.device] | None = None
 
     def _extend(self, psi: torch.Tensor, basic_configs: torch.Tensor | None = None) -> None:
         if basic_configs is None:
@@ -45,7 +54,7 @@ def _extend(self, psi: torch.Tensor, basic_configs: torch.Tensor | None = None)
         logging.info("Number of core configurations: %d", count_core)
 
         self.configs = torch.cat(
-            [self.configs, self.model.find_relative(basic_configs, psi, self.count_extend, self.configs)]
+            [self.configs, self.model.find_relative(basic_configs, psi, self.count_extend, self.configs, self.devices)]
         )
         count_selected = len(self.configs)
         self.psi = torch.nn.functional.pad(self.psi, (0, count_selected - count_core))
@@ -80,7 +89,7 @@ def package(
                 selected = (psi.conj() * psi).real.argsort(descending=True)[: self.count_extend]
                 configs = self.configs
                 self._extend(psi[selected], self.configs[selected])
-                psi = self.model.apply_within(configs, psi, self.configs)
+                psi = self._apply_within(configs, psi, self.configs)
             for _, [alpha, beta, v] in zip(range(1 + self.step), self._run()):
                 yield package(self._eigh_tridiagonal(alpha, beta, v))
         elif self.single_extend:
@@ -181,6 +190,7 @@ def _apply_within(self, configs_i: torch.Tensor, psi_i: torch.Tensor, configs_j:
                 configs_i[start_index:end_index],
                 psi_i[start_index:end_index],
                 configs_j,
+                self.devices,
             )
             if result is None:
                 result = local_result
@@ -429,6 +439,7 @@ def main(
             target_energy: torch.Tensor
             lanczos_results: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
 
+            devices = _get_devices_from_context(context)
             for lanczos_results in _DynamicLanczos(
                 model=model,
                 configs=configs,
@@ -441,6 +452,7 @@ def main(
                 first_extend=self.krylov_extend_first,
                 eigen_count=self.krylov_eigen_count,
                 period=self.krylov_period,
+                devices=devices,
             ).run():
                 target_energy, configs, original_psi = lanczos_results[0]
                 logging.info(
@@ -546,7 +558,7 @@ def closure() -> torch.Tensor:
 
             loss = typing.cast(torch.Tensor, torch.enable_grad(closure)())  # type: ignore[no-untyped-call,call-arg]
             psi: torch.Tensor = loss.psi  # type: ignore[attr-defined]
-            final_energy = ((psi.conj() @ model.apply_within(configs, psi, configs)) / (psi.conj() @ psi)).real
+            final_energy = ((psi.conj() @ model.apply_within(configs, psi, configs, devices)) / (psi.conj() @ psi)).real
             logging.info(
                 "Loss during local optimization: %.10f, Final energy: %.10f, Target energy: %.10f, Reference energy: %.10f, Final error: %.10f",
                 loss.item(),

qmp/algorithms/haar_with_orbit.py

@@ -23,6 +23,7 @@
     _DynamicLanczos,
     _sampling_from_last_iteration,
     _merge_pool_from_neural_network_and_pool_from_last_iteration,
+    _get_devices_from_context,
 )
 from .orbit import NaturalOrbitCalculator, _read_fcidump_tensors
 from ..models.optimized_basis import Model as OptimizedModel, ModelConfig as OptimizedModelConfig
@@ -281,6 +282,8 @@ def main(
             target_energy: torch.Tensor
             lanczos_results: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
 
+            # Get devices for multi-GPU computation
+            devices = _get_devices_from_context(context)
             for lanczos_results in _DynamicLanczos(
                 model=model,
                 configs=configs,
@@ -293,6 +296,7 @@ def main(
                 first_extend=self.krylov_extend_first,
                 eigen_count=self.krylov_eigen_count,
                 period=self.krylov_period,
+                devices=devices,
             ).run():
                 target_energy, configs, original_psi = lanczos_results[0]
                 logging.info("Current energy: %.10f, samples: %d", target_energy.item(), len(configs))
@@ -393,7 +397,7 @@ def closure() -> torch.Tensor:
 
             loss = typing.cast(torch.Tensor, torch.enable_grad(closure)())  # type: ignore[no-untyped-call,call-arg]
             psi: torch.Tensor = loss.psi  # type: ignore[attr-defined]
-            final_energy = ((psi.conj() @ model.apply_within(configs, psi, configs)) / (psi.conj() @ psi)).real
+            final_energy = ((psi.conj() @ model.apply_within(configs, psi, configs, devices)) / (psi.conj() @ psi)).real
             logging.info(
                 "Loss: %.10f, Final energy: %.10f, Target energy: %.10f, Reference energy: %.10f, Final error: %.10f",
                 loss.item(),

qmp/algorithms/pert.py

@@ -6,10 +6,19 @@
 import typing
 import dataclasses
 import omegaconf
+import torch
 from ..utility.context import RuntimeContext
 from ..utility.action_dict import action_dict
 
 
+def _get_devices_from_context(context: RuntimeContext) -> list[torch.device]:
+    """
+    Get the devices list from the runtime context.
+    """
+    assert context.devices is not None
+    return context.devices
+
+
 @dataclasses.dataclass
 class PerturbationConfig:
     """
@@ -29,13 +38,16 @@ def main(
         model = context.create_model(runtime_config.model)
         data = checkpoint_data
 
+        # Get devices for multi-GPU computation
+        devices = _get_devices_from_context(context)
+
         if "haar" not in data and "imag" in data:
             data["haar"] = data.pop("imag")
         configs, psi = data["haar"]["pool"]
         configs = configs.to(context.device)
         psi = psi.to(context.device)
 
-        energy0_num = psi.conj() @ model.apply_within(configs, psi, configs)
+        energy0_num = psi.conj() @ model.apply_within(configs, psi, configs, devices)
         energy0_den = psi.conj() @ psi
         energy0 = (energy0_num / energy0_den).real.item()
         logging.info("Current energy is %.8f", energy0)
@@ -46,7 +58,7 @@ def main(
         current_target_number = number
         logging.info("Starting finding relative configurations with %d.", number)
         while True:
-            other_configs = model.find_relative(configs, psi, current_target_number, configs)
+            other_configs = model.find_relative(configs, psi, current_target_number, configs, devices)
             current_result_number = other_configs.size(0)
             logging.info("Found %d relative configurations.", current_result_number)
             if current_result_number == last_result_number:
@@ -56,9 +68,9 @@ def main(
             logging.info("Doubling target number to %d.", current_target_number)
             break
 
-        hamiltonian_psi = model.apply_within(configs, psi, other_configs)
+        hamiltonian_psi = model.apply_within(configs, psi, other_configs, devices)
         energy2_num = (hamiltonian_psi.conj() * hamiltonian_psi).real / (psi.conj() @ psi).real
-        energy2_den = energy0 - model.diagonal_term(other_configs).real
+        energy2_den = energy0 - model.diagonal_term(other_configs, devices).real
         energy2 = (energy2_num / energy2_den).sum().item()
         logging.info("Correct energy is %.8f", energy2)
         logging.info("Error is reduced from %.8f to %.8f", energy0 - model.ref_energy, energy2 - model.ref_energy)

qmp/algorithms/vmc.py

@@ -12,6 +12,14 @@
 from ..utility.action_dict import action_dict
 
 
+def _get_devices_from_context(context: RuntimeContext) -> list[torch.device]:
+    """
+    Get the devices list from the runtime context.
+    """
+    assert context.devices is not None
+    return context.devices
+
+
 @dataclasses.dataclass
 class VmcConfig:
     """
@@ -47,6 +55,9 @@ def main(
             runtime_config.optimizer, network.parameters(), checkpoint_data.get("optimizer")
         )
 
+        # Get devices for multi-GPU computation
+        devices = _get_devices_from_context(context)
+
         logging.info(
             "Arguments Summary: Sampling Count: %d, Relative Count: %d, Local Steps: %d, Unique: %s",
             self.sampling_count,
@@ -79,7 +90,7 @@ def main(
             else:
                 configs_src = configs_i
                 configs_dst = torch.cat(
-                    [configs_i, model.find_relative(configs_i, psi_i, self.relative_count - len(configs_i))]
+                    [configs_i, model.find_relative(configs_i, psi_i, self.relative_count - len(configs_i), devices=devices)]
                 )
             logging.info("Relative configurations calculated, count: %d", len(configs_dst))
 
@@ -89,7 +100,7 @@ def closure() -> torch.Tensor:
                 psi_src = network(configs_src)
                 with torch.no_grad():
                     psi_dst = network(configs_dst)
-                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src)
+                    hamiltonian_psi_dst = model.apply_within(configs_dst, psi_dst, configs_src, devices)
                 psi_src_reweight = psi_src.conj() * reweight
                 num = psi_src_reweight @ hamiltonian_psi_dst
                 den = psi_src_reweight @ psi_src.detach()