Fix load_accelerator_state only restoring one RNG backend

[Rakshit-gen]

Fixes #3960

save_accelerator_state uses independent if statements for each RNG backend, so it saves a seed for every backend that reports as available. On the load side, HPU, Neuron and CUDA were chained together with elif, ending in a bare else for CUDA. That causes two separate problems.

If more than one backend reports as available, only the first one in the elif chain actually gets its RNG state restored on load_state, so training resumes with a stale RNG state on the others. That's the reproducibility issue reported in #3960.

There's a second problem the elif/else structure causes that wasn't in the original report: if none of XPU, MLU, SDAA, MUSA, HPU, Neuron or CUDA are available, for example a plain CPU machine or a TPU/XLA setup, the else branch still runs unconditionally and tries to read a torch_cuda_manual_seed key that save_accelerator_state never wrote (since it only writes that key when is_cuda_available() is true). That raises a KeyError, which gets silently swallowed by the surrounding try/except, and the xm.set_rng_state call that comes right after in the same block never runs, so on XLA the saved XLA RNG state never actually gets restored on resume.

Change

Switched HPU, Neuron and CUDA in load_accelerator_state from an elif chain to independent if statements, matching the structure save_accelerator_state already uses. Each backend that was actually saved gets restored now, regardless of what else is or isn't available. MLU/SDAA/MUSA stay as an elif chain since save_accelerator_state treats them the same way.

#3960 already reported the first half of this and #3988 proposed close to the same fix, but it went stale without a maintainer review and got closed along with the issue. The underlying bug is still on main, so picking that back up here with the same approach, plus a couple of regression tests that were missing before.

Tests

Added two tests to tests/test_state_checkpointing.py:

• test_rng_state_loads_without_accelerator_specific_backend: patches every specific backend to unavailable and asserts loading a saved checkpoint logs "All random states loaded successfully" instead of silently failing.
• test_rng_state_loads_for_every_saved_backend: patches two backends (HPU and CUDA) as available at once and asserts both get their RNG state restored, not just the first one in the chain.

Checked both against the code before this change, they fail there and pass after. Also ran the full tests/test_state_checkpointing.py suite (CPU only, 22 tests), all pass.

[Rakshit-gen]

@SunMarc could you take a look when you get a chance?