From 49d5ea49bc351d41ef4a87706117fc7f021b8497 Mon Sep 17 00:00:00 2001 From: sunnydave234 <73565162+sunnydave234@users.noreply.github.com> Date: Thu, 30 Jul 2026 03:22:13 -0700 Subject: [PATCH] fix(utils): add MPS branch to torch RNG state serialization (#4014) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit serialize_torch_rng_state/deserialize_torch_rng_state only handled CPU and CUDA generators. On MPS, resumed training was not bit-exact for any stochastic op (dropout, ACT's CVAE noise) since the MPS generator's state was never saved or restored. Mirrors the existing CUDA branch using torch.mps.get_rng_state/set_rng_state (available since torch 2.11). Note: get_rng_state()/set_rng_state() (used by seeded_context()) have the same gap but are out of scope here — happy to follow up separately if useful. Co-authored-by: Sunny Dave Co-authored-by: Steven Palma --- src/lerobot/utils/random_utils.py | 4 ++++ tests/utils/test_random_utils.py | 11 +++++++++++ 2 files changed, 15 insertions(+) diff --git a/src/lerobot/utils/random_utils.py b/src/lerobot/utils/random_utils.py index e280fc342..6edcf2155 100644 --- a/src/lerobot/utils/random_utils.py +++ b/src/lerobot/utils/random_utils.py @@ -85,6 +85,8 @@ def serialize_torch_rng_state() -> dict[str, torch.Tensor]: torch_rng_state_dict = {"torch_rng_state": torch.get_rng_state()} if torch.cuda.is_available(): torch_rng_state_dict["torch_cuda_rng_state"] = torch.cuda.get_rng_state() + if torch.backends.mps.is_available(): + torch_rng_state_dict["torch_mps_rng_state"] = torch.mps.get_rng_state() return torch_rng_state_dict @@ -95,6 +97,8 @@ def deserialize_torch_rng_state(rng_state_dict: dict[str, torch.Tensor]) -> None torch.set_rng_state(rng_state_dict["torch_rng_state"]) if torch.cuda.is_available() and "torch_cuda_rng_state" in rng_state_dict: torch.cuda.set_rng_state(rng_state_dict["torch_cuda_rng_state"]) + if torch.backends.mps.is_available() and "torch_mps_rng_state" in rng_state_dict: + torch.mps.set_rng_state(rng_state_dict["torch_mps_rng_state"]) def serialize_rng_state() -> dict[str, torch.Tensor]: diff --git a/tests/utils/test_random_utils.py b/tests/utils/test_random_utils.py index e3a5d420f..c79e3d506 100644 --- a/tests/utils/test_random_utils.py +++ b/tests/utils/test_random_utils.py @@ -73,6 +73,17 @@ def test_serialize_deserialize_torch_rng(fixed_seed): assert val2 == val3 +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available") +def test_serialize_deserialize_torch_rng_mps(fixed_seed): + _ = torch.rand(1, device="mps").item() + st = serialize_torch_rng_state() + assert "torch_mps_rng_state" in st + val2 = torch.rand(1, device="mps").item() + deserialize_torch_rng_state(st) + val3 = torch.rand(1, device="mps").item() + assert val2 == val3 + + def test_serialize_deserialize_rng(fixed_seed): # Generate one from each library _ = random.random()