Skip to content

Check devices matching during load random seed for device. #48

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hzhangxyz merged 1 commit into from
Jun 19, 2025
Merged

Check devices matching during load random seed for device. #48

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions qmb/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}

Copilot AI Jun 18, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Storing only the device type may lead to mismatches on multi-GPU setups (e.g., cuda:0 vs cuda:1). Consider saving the full device spec (such as str(self.device) or the device index) to ensure the random state is loaded onto the correct device.

Suggested change
data["random"] = {"host": torch.get_rng_state(), "device": dump_random_engine_state(self.device), "device_type": self.device.type}
data["random"] = {
"host": torch.get_rng_state(),
"device": dump_random_engine_state(self.device),
"device_spec": str(self.device)
}

Copilot uses AI. Check for mistakes.

@BrevityD BrevityD Jun 19, 2025

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

多GPU情况下会出现问题吗？我怎么觉得不会？

data_pth = self.folder() / "data.pth"
local_data_pth = self.folder() / f"data.{step}.pth"
torch.save(data, local_data_pth)
Expand Down Expand Up @@ -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:

Copilot AI Jun 18, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comparing only the device type (e.g., "cuda") may miss mismatches in specific device indices (e.g., "cuda:0" vs "cuda:1"). Consider saving and checking the full device identifier via str(self.device) or self.device.index.

Suggested change
if data["random"]["device_type"] == self.device.type:
if data["random"]["device"] == str(self.device):

Copilot uses AI. Check for mistakes.
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")

Copilot AI Jun 18, 2025

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nitpick] This skip notice might warrant a logging.warning (or a lower debug level) instead of info, to better surface potential misconfiguration.

Suggested change
logging.info("Skipping loading random engine state for device since the device type does not match")
logging.warning("Skipping loading random engine state for device since the device type does not match")

Copilot uses AI. Check for mistakes.
else:
logging.info("Random seed not specified, using current seed: %d", torch.seed())

Expand Down
2 changes: 0 additions & 2 deletions qmb/imag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down
2 changes: 0 additions & 2 deletions qmb/rldiag.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]:
Expand Down Expand Up @@ -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")

Expand Down
2 changes: 0 additions & 2 deletions qmb/vmc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down