[device_mesh] improve device selection logic #150897

wanchaol · 2025-04-09T03:13:36Z

Stack from ghstack (oldest at bottom):

as titled, this PR improves the device selection logic when user did not
set the device before calling the DeviceMesh constructor, as a device
manager, DeviceMesh should try to set the device for users in a good
way.

The behavior of set_device before:

If user call init_process_group to init a world process group, we assume user already called set_device and we don't set the device for the user
If user does not init a world process group by themselves, we init a world process group for the user and follow a heuristic to set the device.
This is ok but sometimes the set_device heuristic wouldn't work well (i.e. if user use TORCH_CUDA_VISBILE_DEVICES

So this PR improves the device selection logic to:

If user call init_process_group to init a world process group and the default cuda context is initialized, then we assume user must called some cuda operation before therefore must have selected the device by themselves
If not the above, then we check if envvars have "LOCAL_RANK" and "WORLD_SIZE" from the launcher (i.e. torchrun), if so, we use "LOCAL_RANK" to set the device for the current process, which is a very standard practice. (This solves the TORCH_CUDA_VISBILE_DEVICES issue)
If not above, then we fallback to the old heuristic.

cc @H-Huang @awgu @fegin @fduwjj @wz337 @wconstab @d4l3k @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @jerryzh168

[ghstack-poisoned]

pytorch-bot · 2025-04-09T03:13:39Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/150897

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 0f1daaa with merge base 843e4d1 ():

NEW FAILURE - The following job has failed:

Lint / lintrunner-noclang / linux-job (gh)
>>> Lint for torch/_inductor/config.py:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[ghstack-poisoned]

as titled, this PR improves the device selection logic when user did not set the device before calling the DeviceMesh constructor, as a device manager, DeviceMesh should try to set the device for users in a good way. ghstack-source-id: 8d27c0d Pull Request resolved: #150897

[ghstack-poisoned]

as titled, this PR improves the device selection logic when user did not set the device before calling the DeviceMesh constructor, as a device manager, DeviceMesh should try to set the device for users in a good way. ghstack-source-id: 7967a39 Pull Request resolved: #150897

[ghstack-poisoned]

as titled, this PR improves the device selection logic when user did not set the device before calling the DeviceMesh constructor, as a device manager, DeviceMesh should try to set the device for users in a good way. ghstack-source-id: 55e85d1 Pull Request resolved: #150897

tianyu-l

sorry not having enough context on DeviceMesh, so asking some questions before I can review. Meanwhile @fegin if he could unblock.

tianyu-l · 2025-04-29T03:56:06Z

torch/distributed/device_mesh.py

-                        f"{world_size} ranks and {num_devices_per_host} {self.device_type} devices!"
-                    )
-                device_handle.set_device(get_rank() % num_devices_per_host)
+            if device_handle and not device_handle.is_initialized():


I'm not familiar with the code here. Could you elaborate a bit what are being done here?
It looks like

before this PR, the code would be executed if device_handle exists and c10d is_initialized() is false.

after this PR, the code would be executed if device_handle exists and device_handle's is_initialized() is false.

How is this related to

when user did not set the device before calling the DeviceMesh constructor

This is_initialized() is either torch.cuda.is_initialized() or torch.cpu.is_initialized(). For CUDA, this means the device state is initialized, iiuc.

Yes, I also updated the PR summary to reflect the detailed changes this PR did.

according to your PR desc, we don't do this set_device

If user call init_process_group to init a world process group and the default cuda context is initialized

To translate it into this condition, wouldn't it be?

Suggested change

if device_handle and not device_handle.is_initialized():

if not default_initialized or (device_handle and not device_handle.is_initialized()):

tianyu-l · 2025-04-29T03:57:26Z

torch/distributed/device_mesh.py

+                    # NOTE: This device selection would only work for homogeneous hardware.
+                    num_devices_per_host = device_handle.device_count()
+                    if (
+                        world_size > num_devices_per_host


Could you educate on the relationship between world_size and device_handle vs. global is_initialized call

device_handle is just for device agnostic code (i.e. user might have cuda/xpu/mtia). world_size and is_initialized is to check if the default world process group is initialized or not.

fegin · 2025-04-30T06:47:04Z

torch/distributed/device_mesh.py

-        and serves as a proxy for communication among the device lists within the cluster.
+        DeviceMesh could be used to setup the N dimensional device connections across the cluster,
+        and manage the ProcessGroups for N dimensional parallelisms. Communications could happen on
+        each dimension of the DeviceMesh separately. DeviceMesh respect the device that user select


Suggested change

each dimension of the DeviceMesh separately. DeviceMesh respect the device that user select

each dimension of the DeviceMesh separately. DeviceMesh respects the device that user select

fegin · 2025-04-30T06:48:44Z

torch/distributed/device_mesh.py


-        DeviceMesh can be used as a context manager.
+        DeviceMesh can also be used as a context manager when using together with most DTensor APIs.


Curious, what are the APIs that this doesn't work?

iirc every public DTensor API should work, I just want to be conservative here. All non-DTensor APIs would probably not work because the context manager is just stashing the device mesh created to the mesh env, nothing magic

Update

a7fb0e9

[ghstack-poisoned]

pytorch-bot bot added the oncall: distributed Add this issue/PR to distributed oncall triage queue label Apr 9, 2025

This was referenced Apr 9, 2025

Fix DTensorTestBase to barrier with device ids #150896

Closed

[device_mesh] replace dim_group_info with group_name #150898

Open

wanchaol added the release notes: distributed (dtensor) release notes category label Apr 9, 2025

pytorchbot added the open source label Apr 9, 2025

wanchaol added 2 commits April 9, 2025 18:13

Update

24f2ab1

[ghstack-poisoned]

Update

b392624

[ghstack-poisoned]

pytorch-bot bot added the module: cpu CPU specific problem (e.g., perf, algorithm) label Apr 21, 2025

Update

979cd45

[ghstack-poisoned]

Update

0f1daaa

[ghstack-poisoned]

wanchaol requested review from fegin, wconstab, wz337 and tianyu-l April 21, 2025 20:59

wanchaol added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 21, 2025

tianyu-l reviewed Apr 29, 2025

View reviewed changes

fegin reviewed Apr 30, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[device_mesh] improve device selection logic #150897

[device_mesh] improve device selection logic #150897

wanchaol commented Apr 9, 2025 •

edited

Loading

pytorch-bot bot commented Apr 9, 2025 •

edited

Loading

tianyu-l left a comment

tianyu-l Apr 29, 2025

fegin Apr 30, 2025

wanchaol Apr 30, 2025

tianyu-l May 1, 2025

tianyu-l Apr 29, 2025

wanchaol Apr 30, 2025

fegin Apr 30, 2025 •

edited

Loading

fegin Apr 30, 2025

wanchaol Apr 30, 2025

	if device_handle and not device_handle.is_initialized():
	if not default_initialized or (device_handle and not device_handle.is_initialized()):

	each dimension of the DeviceMesh separately. DeviceMesh respect the device that user select
	each dimension of the DeviceMesh separately. DeviceMesh respects the device that user select


		DeviceMesh can be used as a context manager.
		DeviceMesh can also be used as a context manager when using together with most DTensor APIs.

[device_mesh] improve device selection logic #150897

Are you sure you want to change the base?

[device_mesh] improve device selection logic #150897

Conversation

wanchaol commented Apr 9, 2025 • edited Loading

pytorch-bot bot commented Apr 9, 2025 • edited Loading

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/150897

❌ 1 New Failure

tianyu-l left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fegin Apr 30, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

wanchaol commented Apr 9, 2025 •

edited

Loading

pytorch-bot bot commented Apr 9, 2025 •

edited

Loading

fegin Apr 30, 2025 •

edited

Loading