admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-30 13:09:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(utils): add MPS branch to torch RNG state serialization (#4014)

Browse Source 
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 <sunnydave@Sunnys-Mac-Studio.local>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
This commit is contained in:
sunnydave234
2026-07-30 03:22:13 -07:00
committed by GitHub
parent d23b65416f
commit 49d5ea49bc
2 changed files with 15 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/utils/test_random_utils.py
+11
View File
@@ -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()