Replace NVML device affinity with XPUML

conftest.py

@@ -25,16 +25,19 @@ def original_datadir(request: pytest.FixtureRequest) -> Path:
 
 @cache
 def _get_available_gpus() -> int:
-    import pynvml
+    import torch_xmlir._XMLIRC as XMLIR_C
 
+    initialized = False
     try:
-        pynvml.nvmlInit()
-        device_count = pynvml.nvmlDeviceGetCount()
-        pynvml.nvmlShutdown()
-        return device_count
-    except pynvml.NVMLError as e:
-        print(f"WARNING: Failed to get available GPUs: {e}")
+        XMLIR_C.xpumlInit()
+        initialized = True
+        return XMLIR_C.xpumlDeviceGetCount()
+    except Exception as e:
+        print(f"WARNING: Failed to get available devices: {e}")
         return 0
+    finally:
+        if initialized:
+            XMLIR_C.xpumlShutdown()
 
 
 def pytest_addoption(parser: pytest.Parser):

cosmos_framework/callbacks/device_monitor.py

@@ -6,8 +6,8 @@
 
 import pandas as pd
 import psutil
-import pynvml
 import torch
+import torch_xmlir._XMLIRC as XMLIR_C
 import wandb
 
 from cosmos_framework.callbacks.every_n import EveryN
@@ -102,7 +102,14 @@ def on_train_start(self, model, iteration=0):
             log.info(f"{self.name} callback: local_dir: {self.local_dir}")
 
         local_rank = int(os.getenv("LOCAL_RANK", 0))
-        self.handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank)
+        self._init_device_handle(local_rank)
+
+    def _init_device_handle(self, local_rank: int) -> None:
+        XMLIR_C.xpumlInit()
+        self.handle = XMLIR_C.xpumlDeviceGetHandleByIndex(local_rank)
+
+    def _get_device_memory_info(self) -> Any:
+        return XMLIR_C.xpumlDeviceGetMemoryInfo(self.handle)
 
     def every_n_impl(
         self,
@@ -133,7 +140,7 @@ def every_n_impl(
         util = torch.cuda.utilization()
         clock = torch.cuda.clock_rate()
 
-        memory_info = pynvml.nvmlDeviceGetMemoryInfo(self.handle)
+        memory_info = self._get_device_memory_info()
         nvml_used_gpu_mem_gb = memory_info.used / (1024**3)
         nvml_free_gpu_mem_gb = memory_info.free / (1024**3)
 

cosmos_framework/inference/args.py

@@ -9,7 +9,7 @@
 from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, Self, cast, override
 
 import pydantic
-import pynvml
+import torch_xmlir._XMLIRC as XMLIR_C
 from typing_extensions import assert_never
 from tyro.conf import Suppress
 
@@ -1289,28 +1289,24 @@ def _get_dp_shard_size(
 @cache
 def _get_device_memory_bytes() -> int:
     try:
-        pynvml.nvmlInit()
-        handle = pynvml.nvmlDeviceGetHandleByIndex(0)
-        info = _get_nvml_device_memory_info(handle)
+        XMLIR_C.xpumlInit()
+        handle = XMLIR_C.xpumlDeviceGetHandleByIndex(0)
+        info = _get_xpuml_device_memory_info(handle)
         return info.total
     except Exception:
-        # Fallback for unified memory architectures (e.g., GB10) where
-        # nvmlDeviceGetMemoryInfo is not supported.
+        # Fallback for unified memory architectures where XPUML memory queries
+        # are not supported.
         import torch
+
         if torch.cuda.is_available():
             return int(torch.cuda.get_device_properties(0).total_memory)
         return 128 * 1024**3  # Default 128GB
     finally:
         try:
-            pynvml.nvmlShutdown()
+            XMLIR_C.xpumlShutdown()
         except Exception:
             pass
 
 
-def _get_nvml_device_memory_info(handle: Any) -> Any:
-    try:
-        return pynvml.nvmlDeviceGetMemoryInfo_v2(handle)
-    except AttributeError:
-        return pynvml.nvmlDeviceGetMemoryInfo(handle)
-    except pynvml.NVMLError_NotSupported:
-        return pynvml.nvmlDeviceGetMemoryInfo(handle)
+def _get_xpuml_device_memory_info(handle: Any) -> Any:
+    return XMLIR_C.xpumlDeviceGetMemoryInfo(handle)

cosmos_framework/utils/device.py

@@ -2,11 +2,10 @@
 # SPDX-License-Identifier: OpenMDW-1.1
 
 import gc
-import math
 import os
 from functools import wraps
 
-import pynvml
+import torch_xmlir._XMLIRC as XMLIR_C
 from loguru import logger as logging
 
 
@@ -18,28 +17,28 @@ def get_gpu_architecture():
         str: The GPU architecture, which can be "H100", "A100", or "Other".
     """
     try:
-        pynvml.nvmlInit()
-        device_count = pynvml.nvmlDeviceGetCount()
+        XMLIR_C.xpumlInit()
+        device_count = XMLIR_C.xpumlDeviceGetCount()
         for i in range(device_count):
-            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
-            model_name = pynvml.nvmlDeviceGetName(handle)
+            handle = XMLIR_C.xpumlDeviceGetHandleByIndex(i)
+            model_name = XMLIR_C.xpumlDeviceGetName(handle)
             if isinstance(model_name, bytes):
                 model_name = model_name.decode("utf-8")
-            print(f"GPU {i}: Model: {model_name}")
-
-            # Check for specific models like H100 or A100
-            if "H100" in model_name or "H200" in model_name:
+            if "H100" in model_name:
                 return "H100"
             elif "A100" in model_name:
                 return "A100"
             elif "L40S" in model_name:
                 return "L40S"
             elif "B200" in model_name:
                 return "B200"
-    except pynvml.NVMLError as error:
-        print(f"Failed to get GPU info: {error}")
+            else:
+                return "Other"
+    except Exception as error:
+        logging.error(f"Error retrieving device information: {error}")
+        return "Other"
     finally:
-        pynvml.nvmlShutdown()
+        XMLIR_C.xpumlShutdown()
 
     # return "Other" incase of non hopper/ampere or error
     return "Other"
@@ -55,13 +54,13 @@ class GPUArchitectureNotSupported(Exception):
 
 def print_gpu_mem(str=None):
     try:
-        pynvml.nvmlInit()
-        meminfo = pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0))
+        XMLIR_C.xpumlInit()
+        meminfo = XMLIR_C.xpumlDeviceGetMemoryInfo(XMLIR_C.xpumlDeviceGetHandleByIndex(0))
         logging.info(
             f"{str}: {meminfo.used / 1024 / 1024}/{meminfo.total / 1024 / 1024}MiB used ({meminfo.free / 1024 / 1024}MiB free)"
         )
-    except pynvml.NVMLError as error:
-        print(f"Failed to get GPU memory info: {error}")
+    except Exception as error:
+        print(f"Failed to get device memory info: {error}")
 
 
 def force_gc():
@@ -76,32 +75,37 @@ def force_gc():
 
 def gpu0_has_80gb_or_less():
     try:
-        pynvml.nvmlInit()
-        meminfo = pynvml.nvmlDeviceGetMemoryInfo(pynvml.nvmlDeviceGetHandleByIndex(0))
+        XMLIR_C.xpumlInit()
+        meminfo = XMLIR_C.xpumlDeviceGetMemoryInfo(XMLIR_C.xpumlDeviceGetHandleByIndex(0))
         return meminfo.total / 1024 / 1024 / 1024 <= 80
-    except pynvml.NVMLError as error:
-        print(f"Failed to get GPU memory info: {error}")
+    except Exception as error:
+        print(f"Failed to get device memory info: {error}")
 
 
 class Device:
 
-    _nvml_affinity_elements = math.ceil(os.cpu_count() / 64)  # type: ignore
-
     def __init__(self, device_idx: int):
         super().__init__()
-        self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx)
+        self.device_idx = device_idx
+        XMLIR_C.xpumlInit()
+        self.handle = XMLIR_C.xpumlDeviceGetHandleByIndex(device_idx)
 
     def get_name(self) -> str:
-        return pynvml.nvmlDeviceGetName(self.handle)
+        return XMLIR_C.xpumlDeviceGetName(self.handle)
 
     def get_cpu_affinity(self) -> list[int]:
-        affinity_string = ""
-        for j in pynvml.nvmlDeviceGetCpuAffinity(self.handle, Device._nvml_affinity_elements):
-            # assume nvml returns list of 64 bit ints
-            affinity_string = "{:064b}".format(j) + affinity_string
-        affinity_list = [int(x) for x in affinity_string]
-        affinity_list.reverse()  # so core 0 is in 0th element of list
-        return [i for i, e in enumerate(affinity_list) if e != 0]
+        cpu_count = os.cpu_count() or 1
+        device_count = XMLIR_C.xpumlDeviceGetCount()
+        if device_count <= 0 or cpu_count <= device_count:
+            return list(range(cpu_count))
+
+        cores_per_device = max(1, cpu_count // device_count)
+        start_core = self.device_idx * cores_per_device
+        if self.device_idx == device_count - 1:
+            end_core = cpu_count
+        else:
+            end_core = min(start_core + cores_per_device, cpu_count)
+        return list(range(start_core, end_core))
 
 
 def with_torch_device(device):

cosmos_framework/utils/distributed.py

@@ -12,8 +12,8 @@
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any, Callable, Container, Optional
 
-import pynvml
 import torch
+import torch_xmlir._XMLIRC as XMLIR_C
 import torch.distributed as dist
 from torch.distributed import get_process_group_ranks
 
@@ -32,37 +32,20 @@
 
 def init() -> int | None:
     """Initialize distributed training."""
-    #if dist.is_initialized():
-    #    return torch.cuda.current_device()
-
-    # Set GPU affinity.
-    #pynvml.nvmlInit()
+    # if dist.is_initialized():
+    #     return torch.cuda.current_device()
+    XMLIR_C.xpumlInit()
     local_rank = int(os.getenv("LOCAL_RANK", 0))
-    #try:
-    #    device = Device(local_rank)
-    #    os.sched_setaffinity(0, device.get_cpu_affinity())
-    #except pynvml.NVMLError as e:
-    #    log.warning(f"Failed to set device affinity: {e}")
-    # Set up NCCL communication.
-    if hasattr(os, 'sched_setaffinity'):
-        try:
-        # 获取系统CPU核心数
-            cpu_count = os.cpu_count() or 1
-            # 假设GPU数量
-            num_gpus = torch.cuda.device_count()
-
-            if num_gpus > 0 and cpu_count > num_gpus:
-                # 为每个GPU分配CPU核心
-                cores_per_gpu = cpu_count // num_gpus
-                start_core = local_rank * cores_per_gpu
-                end_core = min(start_core + cores_per_gpu, cpu_count)
-
-                # 设置CPU亲和性
-                affinity_mask = set(range(start_core, end_core))
-                os.sched_setaffinity(0, affinity_mask)
-                log.info(f"Set CPU affinity for GPU {local_rank} to cores {affinity_mask}") 
-        except Exception as e:
-                log.warning(f"Failed to set CPU affinity: {e}")
+    # if hasattr(os, "sched_setaffinity"):
+    try:
+        device = Device(local_rank)
+        affinity_mask = set(device.get_cpu_affinity())
+        if affinity_mask:
+            os.sched_setaffinity(0, affinity_mask)
+            log.info(f"Set CPU affinity for device {local_rank} to cores {affinity_mask}")
+    except Exception as e:
+        log.warning(f"Failed to set CPU affinity: {e}")
+
     os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "0"
     os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = "1"
     if dist.is_available():