Check devices matching during load random seed for device.

Author: hzhangxyz

qmb/common.py

@@ -10,7 +10,7 @@
 import torch
 import tyro
 from .model_dict import model_dict, ModelProto, NetworkProto
-from .random_engine import load_random_engine_state
+from .random_engine import dump_random_engine_state, load_random_engine_state
 
 
 @dataclasses.dataclass
@@ -77,6 +77,7 @@ def save(self, data: typing.Any, step: int) -> None:
         """
         Save data to checkpoint.
         """
+        data["random"] = {"host": torch.get_rng_state(), "device": dump_random_engine_state(self.device), "device_type": self.device.type}
         data_pth = self.folder() / "data.pth"
         local_data_pth = self.folder() / f"data.{step}.pth"
         torch.save(data, local_data_pth)
@@ -151,7 +152,10 @@ def main(self, *, model_param: typing.Any = None, network_param: typing.Any = No
         elif "random" in data:
             logging.info("Loading random seed from the checkpoint")
             torch.set_rng_state(data["random"]["host"])
-            load_random_engine_state(data["random"]["device"], self.device)
+            if data["random"]["device_type"] == self.device.type:
+                load_random_engine_state(data["random"]["device"], self.device)
+            else:
+                logging.info("Skipping loading random engine state for device since the device type does not match")
         else:
             logging.info("Random seed not specified, using current seed: %d", torch.seed())
 

qmb/imag.py

@@ -16,7 +16,6 @@
 from .subcommand_dict import subcommand_dict
 from .model_dict import ModelProto
 from .optimizer import initialize_optimizer, scale_learning_rate
-from .random_engine import dump_random_engine_state
 
 
 @dataclasses.dataclass
@@ -513,7 +512,6 @@ def closure() -> torch.Tensor:
             data["imag"]["global"] += 1
             data["network"] = network.state_dict()
             data["optimizer"] = optimizer.state_dict()
-            data["random"] = {"host": torch.get_rng_state(), "device": dump_random_engine_state(self.common.device)}
             self.common.save(data, data["imag"]["global"])
             logging.info("Checkpoint successfully saved")
 

qmb/rldiag.py

@@ -14,7 +14,6 @@
 from .model_dict import ModelProto
 from .optimizer import initialize_optimizer
 from .bitspack import pack_int
-from .random_engine import dump_random_engine_state
 
 
 def lanczos_energy(model: ModelProto, configs: torch.Tensor, step: int, threshold: float) -> tuple[float, torch.Tensor]:
@@ -215,7 +214,6 @@ def main(self) -> None:
             data["rldiag"]["local"] += 1
             data["network"] = network.state_dict()
             data["optimizer"] = optimizer.state_dict()
-            data["random"] = {"host": torch.get_rng_state(), "device": dump_random_engine_state(self.common.device)}
             self.common.save(data, data["rldiag"]["global"])
             logging.info("Checkpoint successfully saved")
 

qmb/vmc.py

@@ -11,7 +11,6 @@
 from .common import CommonConfig
 from .subcommand_dict import subcommand_dict
 from .optimizer import initialize_optimizer
-from .random_engine import dump_random_engine_state
 
 
 @dataclasses.dataclass
@@ -133,7 +132,6 @@ def closure() -> torch.Tensor:
             data["vmc"]["global"] += 1
             data["network"] = network.state_dict()
             data["optimizer"] = optimizer.state_dict()
-            data["random"] = {"host": torch.get_rng_state(), "device": dump_random_engine_state(self.common.device)}
             self.common.save(data, data["vmc"]["global"])
             logging.info("Checkpoint successfully saved")