[YHC66]

Description

deterministic_rng (added to transformer_lens/benchmarks/utils.py in #1934) forks the CUDA generators with

with torch.random.fork_rng(devices=cuda_devices):

Since torch 2.13, fork_rng no longer defaults device_type to "cuda": it uses torch.accelerator.current_accelerator() and falls back to CUDA only when there is none. Where the current accelerator is not CUDA (MPS on Apple Silicon), the CUDA device list is snapshotted and restored through torch.mps instead, so the CUDA generators are not restored after the block.

On real hardware this is harmless today, because a Mac has no CUDA devices and the list is empty. It does make tests/unit/benchmarks/test_device_generators_are_seeded_inside_and_restored_after fail on macOS with torch >= 2.13 (the device_generators fixture fakes both CUDA and MPS):

AssertionError: assert {'cuda': 7, 'mps': 'caller'} == {'cuda': 'caller', 'mps': 'caller'}

CI does not hit this: the unit tests run on Linux runners with no accelerator (so fork_rng falls back to CUDA), and the macOS job uses the locked torch 2.11.

Changes:

• pass device_type="cuda" to fork_rng, so the CUDA fork does not depend on torch's default (the argument has existed since well before torch 2.6);
• make the device_generators fixture also report MPS as the current accelerator, as on a Mac, so the test covers this on every runner.

test_deterministic_rng.py	torch 2.14.1 (macOS)	torch 2.11.0 (locked)
dev 34c5d77f	1 failed, 1 passed	2 passed
with the fixture change only	1 failed, 1 passed	2 passed
with the fix	2 passed	2 passed

With torch 2.11 fork_rng does not consult the accelerator, so the fixture change is a no-op there.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

No new test: the existing test_device_generators_are_seeded_inside_and_restored_after already covers this once its fixture reports MPS as the current accelerator.

Local runs: tests/unit -m "not slow" on macOS with torch 2.14.1 and transformers 5.20.2: 7651 passed, 0 failed; tests/unit/benchmarks on both torch 2.14.1 and 2.11.0 (101 passed each), plus black, isort, pycln and mypy ..