Probe the device module for train_cifar's fork_rng device entries - #8407
Probe the device module for train_cifar's fork_rng device entries#8407delock wants to merge 1 commit into
Conversation
train_cifar always listed the current device in fork_rng's devices=, but torch.cpu has no get_rng_state and fork_rng already saves the CPU/global RNG. Probe the device module for a per-device RNG instead of matching the accelerator name, so backends without get_rng_state (cpu) pass no device entries at all. The sibling fixes from the original commits (DDP device_ids pinning, device names, fp16 skips) landed separately as deepspeedai#8397, deepspeedai#8398 and deepspeedai#8399. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ade99f4704
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| with get_accelerator().random().fork_rng(devices=[get_accelerator().current_device_name()], **fork_kwargs): | ||
| # fork_rng only needs entries for backends with per-device generators: the global | ||
| # CPU RNG is always saved, and torch.cpu has no get_rng_state to call anyway. | ||
| device_mod = torch.get_device_module(get_accelerator().device_name()) |
There was a problem hiding this comment.
When these tests run with the supported minimum PyTorch 2.0 (requirements/requirements.txt specifies torch>=2.0.0), torch.get_device_module does not exist, so every train_cifar call now raises AttributeError before entering fork_rng, including CUDA paths that previously worked. Please use a lookup available on older supported releases or guard this API by PyTorch version.
Useful? React with 👍 / 👎.
| else: | ||
| fork_kwargs = {} | ||
| with get_accelerator().random().fork_rng(devices=[get_accelerator().current_device_name()], **fork_kwargs): | ||
| # fork_rng only needs entries for backends with per-device generators: the global |
There was a problem hiding this comment.
This is a non-merge commit, but its message has no Signed-off-by trailer. Recreate the commit with --signoff so it satisfies the repository's commit requirements.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
1 participant
Description
train_cifarunconditionally passeddevices=[get_accelerator().current_device_name()]tofork_rng. On CPU that becomesdevices=['cpu'], andtorch.cpuhas noget_rng_state, so the context manager raisesAttributeErrorbefore the test body starts — the CPU RNG lives in the global generator thatfork_rngalready forks.Probe the device module for a per-device RNG instead of matching accelerator names:
get_rng_state(e.g. cuda) behave exactly as before;devices=[], which changes nothing beyond the global-RNG save/restorefork_rngalways performs.The crash is only reachable from the multi-rank tests that call
train_cifar(test_onebit.py,test_pipe.py), which the CPU runner currently skips at the device gate; it was exposed by theLOCAL_SIZE=4experiment in #8381. Sibling fixes from the same series landed as #8397, #8398 and #8399.Test plan
LOCAL_SIZE=4multi-rank CPU run tracked in [DON'T MERGE] Run multi-rank CPU unit tests in CI via LOCAL_SIZE #8381: thetest_onebit/test_pipecallers got pastfork_rngand into their test bodies with this exact change.