fix(mge/dist): fix gl oom error

GitOrigin-RevId: 4ba3d2cfd7

imperative/python/megengine/distributed/helper.py

@@ -186,9 +186,9 @@ def _get_device_count_worker(queue, device_type):
     queue.put(num)
 
 
-def _check_device_initialized(device_type: str):
+def _check_device_initialized(device_type: str, rank: int):
     try:
-        test = Tensor(1, device=device_type)
+        test = Tensor(1, device=(device_type + str(rank)))
         inited = False
         del test
     except:

imperative/python/megengine/distributed/launcher.py

@@ -39,7 +39,7 @@ def _run_wrapped(
     machine_ranks: list,
 ):
     """Init distributed process group and run wrapped function."""
-    _check_device_initialized(device_type)
+    _check_device_initialized(device_type, dev)
     init_process_group(
         master_ip=master_ip,
         port=port,