From 1324ea1bf0bfc7b8d55e6298fe2265e244be43bb Mon Sep 17 00:00:00 2001 From: Yihang Chen Date: Fri, 23 Oct 2026 11:29:17 -0700 Subject: [PATCH] Fork the CUDA RNG explicitly in deterministic_rng Since torch 2.13, torch.random.fork_rng defaults device_type to the current accelerator instead of CUDA. deterministic_rng passes it the CUDA device list, so wherever the current accelerator is not CUDA (MPS on Apple Silicon) that list is snapshotted and restored through the wrong backend. On real hardware the list is empty there, but it makes test_device_generators_are_seeded_inside_and_restored_after fail on macOS. Pass device_type="cuda", and have the test fixture report MPS as the current accelerator so the case is covered on every runner. --- tests/unit/benchmarks/test_deterministic_rng.py | 4 ++++ transformer_lens/benchmarks/utils.py | 3 ++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit/benchmarks/test_deterministic_rng.py b/tests/unit/benchmarks/test_deterministic_rng.py index 9d5e2514c..ef917ca6f 100644 --- a/tests/unit/benchmarks/test_deterministic_rng.py +++ b/tests/unit/benchmarks/test_deterministic_rng.py @@ -21,6 +21,10 @@ def setter(name): monkeypatch.setattr(torch.cuda, "manual_seed", setter("cuda")) monkeypatch.setattr(torch.cuda, "manual_seed_all", setter("cuda")) monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + # As on Apple Silicon, MPS is the current accelerator + monkeypatch.setattr( + torch.accelerator, "current_accelerator", lambda *args, **kwargs: torch.device("mps") + ) monkeypatch.setattr(torch.mps, "get_rng_state", lambda *args, **kwargs: state["mps"]) monkeypatch.setattr(torch.mps, "set_rng_state", setter("mps")) monkeypatch.setattr(torch.mps, "manual_seed", setter("mps")) diff --git a/transformer_lens/benchmarks/utils.py b/transformer_lens/benchmarks/utils.py index a7edcfe00..70de76c9a 100644 --- a/transformer_lens/benchmarks/utils.py +++ b/transformer_lens/benchmarks/utils.py @@ -27,7 +27,8 @@ def deterministic_rng(seed: int = 42) -> Iterator[None]: # generators snapshotted here are seeded, so each one is restored afterwards. cuda_devices = list(range(torch.cuda.device_count())) if torch.cuda.is_available() else [] mps_state = torch.mps.get_rng_state() if torch.backends.mps.is_available() else None - with torch.random.fork_rng(devices=cuda_devices): + # fork_rng defaults to the current accelerator since torch 2.13, so name CUDA explicitly. + with torch.random.fork_rng(devices=cuda_devices, device_type="cuda"): try: torch.default_generator.manual_seed(seed) if cuda_devices: