[device mesh] only check when world size > num_devices per host (#111091

)

as titled
Pull Request resolved: #111091
Approved by: https://github.com/awgu, https://github.com/wz337
ghstack dependencies: #110898, #110900

torch/distributed/_tensor/device_mesh.py

@@ -188,7 +188,10 @@ def _get_or_create_default_group(self):
             # automatically set the current cuda/cuda-like device base on num of gpu devices available in each host
             # 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 != 0:
+            if (
+                world_size > num_devices_per_host
+                and world_size % num_devices_per_host != 0
+            ):
                 raise RuntimeError(
                     f"DeviceMesh only support homogeneous hardware, but found "
                     f"{world_size} ranks and {num_devices_per_host} {self.device_type} devices!"