Use device agnostic APIs for RNG #159021
Use device agnostic APIs for RNG #159021
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159021
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit d9f0ed4 with merge base 6834911 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
There was a problem hiding this comment.
Pull Request Overview
This PR updates test files to use device-agnostic APIs for random number generator (RNG) seeding instead of hardcoded CUDA APIs, making the code more portable across different device types. It also adds proper type annotations to fix MYPY linting issues.
- Replace
torch.cuda.manual_seed()withtorch.get_device_module(self.device_type).manual_seed() - Add comprehensive type annotations to all test methods and helper functions
- Fix a missing manual seed initialization in the distribute region test
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| test/distributed/tensor/test_random_ops.py | Adds type annotations to all methods and adds device-agnostic manual seed call for distribute region test |
| test/distributed/tensor/parallel/test_tp_random_state.py | Replaces CUDA-specific manual_seed with device-agnostic API and adds comprehensive type annotations |
The manual_seed call was earlier placed before dtesnor.uniform_(). Setting the seed before creatring the tensor to ensure reproducible initialization. Signed-off-by: Aby Mathew C <aby.mathew.c@intel.com>
| @@ -130,6 +131,7 @@ def test_meta_tensor_init(self): | |||
| ) | |||
|
|
|||
| # Test 2: disable the distribute region for RNG | |||
| torch.get_device_module(self.device_type).manual_seed(self.rank) | |||
There was a problem hiding this comment.
This change feels suspicious because it adds a new set seed where there previously wasn't one. Also it missed changing a cuda-specific seed setting in the same test further above.
MOTIVATION
In order to make the test files more device agnostic, use device agnostic APIs to set the manual seed rather than hard coded CUDA APIs. We have also verified that this does not break any of the existing CUDA functionality.
CHANGES
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta @EikanWang @kwen2501 @ankurneog