Check devices matching during load random seed for device.

When dumping random state for one device but loading it from another, the format
of random state differs so program will raise error, we need to check whether
the device type is unchanged before loading it.

PR: USTC-KnowledgeComputingLab/qmb#48
Signed-off-by: Hao Zhang <[email protected]>

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")