Check devices matching during load random seed for device.

[hzhangxyz]

Description

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.

Checklist:

• I have read the CONTRIBUTING.md.

[Copilot]

Pull Request Overview

This PR centralizes random engine state handling by removing per-module dumps and adding device-type checks when loading checkpoints.

• Removed redundant dump_random_engine_state calls in vmc.py, rldiag.py, and imag.py.
• Added random state serialization (including device type) in common.save.
• Added a check in common.main to only load device state when the saved device type matches the current one.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
qmb/vmc.py	Removed manual random state dump in the closure (now handled by common.save).
qmb/rldiag.py	Removed manual random state dump in the checkpoint logic (delegated to common).
qmb/imag.py	Removed manual random state dump in the closure (delegated to common).
qmb/common.py	Imported both dump/load functions, added random state saving with device type, and added a conditional load in main.

Comments suppressed due to low confidence (2)

qmb/common.py:155

• The new branch that skips loading when device types differ should be covered by tests to ensure it behaves as expected without raising errors. Please add unit tests simulating mismatched device states.

            if data["random"]["device_type"] == self.device.type:

qmb/common.py:78

• The docstring for save() should be updated to mention that it now also serializes the random engine state fields (host, device, and device_type).

        Save data to checkpoint.

[Copilot]

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

[BrevityD]

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

[BrevityD]

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

[hzhangxyz]

@BrevityD 一开始在cpu上跑然后切换到gpu上继续跑的时候会出问题。