Add TEST_HPU flag to set device type (#153461)

MOTIVATION
This PR includes a minor change to check for TEST_HPU flag as well before falling back to CPU. Without this flag, some tests were falling back to CPU causing them to fail.
Please refer to this RFC as well: pytorch/rfcs#66

CHANGES
add TEST_HPU flag to some of the conditions checking the environment
use DEVICE_COUNT variable instead of torch.accelerator.device_count() API since the later is not supported on out-of-tree devices like Intel Gaudi.
@ankurneog , @EikanWang , @cyyever , @guangyey

Pull Request resolved: #153461
Approved by: https://github.com/EikanWang, https://github.com/cyyever, https://github.com/albanD

Author: amathewc

torch/testing/_internal/distributed/_tensor/common_dtensor.py

@@ -58,7 +58,7 @@
 NUM_DEVICES = 4
 
 # We use this as a proxy for "multiple GPUs exist"
-if (TEST_CUDA or TEST_XPU) and DEVICE_COUNT > 1:
+if (TEST_CUDA or TEST_XPU or TEST_HPU) and DEVICE_COUNT > 1:
     # when we actually have multiple GPUs, relax the requirement to smaller counts.
     NUM_DEVICES = min(NUM_DEVICES, DEVICE_COUNT)
 
@@ -339,11 +339,8 @@ def world_size(self) -> int:
 
     @property
     def device_type(self) -> str:
-        # if enough GPU we can use GPU, otherwise we fallback to CPU
-        if (
-            not (TEST_CUDA or TEST_XPU)
-            or torch.accelerator.device_count() < self.world_size
-        ):
+        # if enough GPU/XPU/HPU we can use those devices, otherwise we fallback to CPU
+        if not (TEST_CUDA or TEST_XPU or TEST_HPU) or DEVICE_COUNT < self.world_size:
             return "cpu"
         else:
             return DEVICE_TYPE