Repository: oumi-ai/oumi

Issue:
[Bug] pynvml not recognizing the GPU name
### What happened?

This was reported in [this discord channel](https://discord.com/channels/1286348126797430814/1291937814661038121/1398974741225668740)

When the run the "oumi train -c configs/recipes/vision/qwen2_5_vl_3b/sft/full/train.yaml" on server with Multi-Instance GPU (MIG) they get the following error:

<img width="942" height="858" alt="Image" src="https://github.com/user-attachments/assets/d4a534f2-d7b4-46e3-b414-ad03e7ade235" />

2025-07-27 15:47:10,211][oumi][rank0][pid: 15669] [MainThread][ERROR]][device_utils.py:223] Failed to get GPU utilization for device: 0
Traceback (most recent call last):
File "/home/aru_sharma/test_notebook/venv _again/lib/python3.11/site-packages/oumi/utils/device_utils.py", line 219, in _get _nvidia_gpu_runtime_info_impl
へへへへへへへへへへへへへへへへへへへへへへへへへへへへへへへへへへへへ
File "/home/aru_sharma/test_notebook/venv_again/lib/python3.11/site-packages/pynvml.py", line 3015, in nvmlDeviceGetUtilizationRates
nvmlCheckReturn (ret)
File "/home/aru_sharma/test_notebook/venv_again/lib/python3.11/site-packages/pynvml.py", line 979, in _nvmlCheckReturn raise NVMLError (ret)
pynvml.NVMLError_NotSupported: Not Supported
[2025-07-27 15:47:10,213][oumz][rank0][pid: 15669][MainThread][INF0]][device_utils.py:297] GPU Metrics Before Training: GPU runtime info: None.
[2025-07-27 15:47:10,213][oumi][rank0][pid: 15669][Ma inThread][INFO]][train.py:508] Training init time: 18.479s
[2025-07-27 15:47:10,213][oumi][rank][pid:15669][MainThread][INFO]][train.py:509] Starting training... (TrainerType. TRL SFT, transformers: 4.51.3)

### Steps to reproduce the bug

1. run oumi train -c configs/recipes/vision/qwen2_5_vl_3b/sft/full/train.yaml on server with Multi-Instance GPU (MIG)

### System Info

```shell
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   Oumi environment information:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌────────────────┬──────────────────────────────────────────────┐
│ Oumi version   │ 0.2.2.dev55+ge68a40b2                        │
│ Python version │ 3.11.9                                       │
│ Platform       │ Linux-6.8.0-64-generic-x86_64-with-glibc2.35 │
└────────────────┴──────────────────────────────────────────────┘

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   Installed dependencies:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┏━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┓
┃ PACKAGE          ┃ VERSION         ┃
┡━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━┩
│ accelerate       │ 1.7.0           │
│ aiohttp          │ 3.11.18         │
│ bitsandbytes     │ 0.45.5          │
│ datasets         │ 3.2.0           │
│ diffusers        │ <not installed> │
│ einops           │ 0.8.1           │
│ jsonlines        │ 4.0.0           │
│ liger-kernel     │ 0.5.10          │
│ llama-cpp-python │ <not installed> │
│ lm-eval          │ 0.4.8           │
│ mlflow           │ 3.1.4           │
│ numpy            │ 1.26.4          │
│ nvidia-ml-py     │ 12.560.30       │
│ omegaconf        │ 2.4.0.dev3      │
│ open_clip_torch  │ <not installed> │
│ pandas           │ 2.3.1           │
│ peft             │ 0.15.2          │
│ pexpect          │ 4.8.0           │
│ pillow           │ 11.1.0          │
│ pydantic         │ 2.9.2           │
│ responses        │ 0.25.7          │
│ sglang           │ <not installed> │
│ skypilot         │ 0.7.0           │
│ tensorboard      │ 2.18.0          │
│ timm             │ <not installed> │
│ torch            │ 2.6.0           │
│ torchdata        │ 0.9.0           │
│ torchvision      │ 0.21.0          │
│ tqdm             │ 4.67.1          │
│ transformers     │ 4.51.3          │
│ trl              │ 0.18.2          │
│ typer            │ 0.16.0          │
│ vllm             │ 0.8.3           │
│ wandb            │ 0.19.11         │
└──────────────────┴─────────────────┘

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   Environment variables:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ VARIABLE                        ┃ VALUE                                    ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
│ ACCELERATE_DYNAMO_BACKEND       │ <not set>                                │
│ ACCELERATE_DYNAMO_MODE          │ <not set>                                │
│ ACCELERATE_DYNAMO_USE_DYNAMIC   │ <not set>                                │
│ ACCELERATE_DYNAMO_USE_FULLGRAPH │ <not set>                                │
│ ACCELERATE_USE_FSDP             │ <not set>                                │
│ CUDA_VISIBLE_DEVICES            │ MIG-88145352-5e1e-5a2a-a3f9-26b2799a0e3c │
│ LOCAL_RANK                      │ <not set>                                │
│ LOCAL_WORLD_SIZE                │ <not set>                                │
│ OUMI_EXTRA_DEPS_FILE            │ <not set>                                │
│ OUMI_FORCE_EDITABLE_INSTALL     │ <not set>                                │
│ OUMI_SLURM_CONNECTIONS          │ <not set>                                │
│ OUMI_USE_SPOT_VM                │ <not set>                                │
│ RANK                            │ <not set>                                │
│ WORLD_SIZE                      │ <not set>                                │
└─────────────────────────────────┴──────────────────────────────────────────┘

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
   PyTorch information:
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

┌────────────────┬─────────────────────────────┐
│ CUDA available │ True                        │
│ CUDA version   │ 12.4                        │
│ cuDNN version  │ 90.1.0                      │
│ Number of GPUs │ 1                           │
│ GPU type       │ NVIDIA H100 NVL MIG 3g.47gb │
│ GPU memory     │ 46.4GB                      │
└────────────────┴─────────────────────────────┘
```


Discussion:
This looks related to pynvml not recognizing the GPU name. Ideally we should ignore instead of raising an error since it's only used for logging / telemetry
Relevant:
https://github.com/pytorch/serve/issues/1722



Patch:
diff --git a/src/oumi/utils/device_utils.py b/src/oumi/utils/device_utils.py
index 5939b3a4fa..22fe9323a4 100644
--- a/src/oumi/utils/device_utils.py
+++ b/src/oumi/utils/device_utils.py
@@ -144,6 +144,10 @@ def _get_nvidia_gpu_runtime_info_impl(
 
     try:
         gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(device_index)
+    except pynvml.NVMLError_NotSupported:
+        # This error is expected on some systems.
+        # Only do DEBUG-level logging to reduce noise.
+        logger.debug(f"pyNVML GPU handle not supported for device: {device_index}")
     except Exception:
         logger.exception(f"Failed to get GPU handle for device: {device_index}")
         return None
@@ -153,6 +157,12 @@ def _get_nvidia_gpu_runtime_info_impl(
         try:
             info = pynvml.nvmlDeviceGetMemoryInfo(gpu_handle)
             used_memory_mb_value = float(info.used) // 1024**2
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU memory info not supported for device: {device_index}"
+            )
         except Exception:
             logger.exception(
                 f"Failed to get GPU memory info for device: {device_index}"
@@ -165,6 +175,12 @@ def _get_nvidia_gpu_runtime_info_impl(
             temperature_value = pynvml.nvmlDeviceGetTemperature(
                 gpu_handle, pynvml.NVML_TEMPERATURE_GPU
             )
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU temperature not supported for device: {device_index}"
+            )
         except Exception:
             logger.exception(
                 f"Failed to get GPU temperature for device: {device_index}"
@@ -176,6 +192,12 @@ def _get_nvidia_gpu_runtime_info_impl(
     if fan_speed:
         try:
             fan_speed_value = pynvml.nvmlDeviceGetFanSpeed(gpu_handle)
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU fan speed not supported for device: {device_index}"
+            )
         except Exception:
             # The `GetFanSpeed` function fails on many systems
             # Only do DEBUG-level logging to reduce noise.
@@ -206,6 +228,12 @@ def _get_nvidia_gpu_runtime_info_impl(
 
             milliwatts = pynvml.nvmlDeviceGetPowerManagementLimit(gpu_handle)
             power_limit_watts_value = float(milliwatts) * 1e-3
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU power usage not supported for device: {device_index}"
+            )
         except Exception:
             logger.exception(
                 f"Failed to get GPU power usage for device: {device_index}"
@@ -219,6 +247,12 @@ def _get_nvidia_gpu_runtime_info_impl(
             result = pynvml.nvmlDeviceGetUtilizationRates(gpu_handle)
             gpu_utilization_value = int(result.gpu)
             memory_utilization_value = int(result.memory)
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU utilization not supported for device: {device_index}"
+            )
         except Exception:
             logger.exception(
                 f"Failed to get GPU utilization for device: {device_index}"
@@ -231,6 +265,12 @@ def _get_nvidia_gpu_runtime_info_impl(
             performance_state_value = int(
                 pynvml.nvmlDeviceGetPerformanceState(gpu_handle)
             )
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU performance state not supported for device: {device_index}"
+            )
         except Exception:
             logger.exception(
                 f"Failed to get GPU performance state for device: {device_index}"
@@ -251,6 +291,12 @@ def _get_nvidia_gpu_runtime_info_impl(
             clock_speed_memory_value = int(
                 pynvml.nvmlDeviceGetClockInfo(gpu_handle, pynvml.NVML_CLOCK_MEM)
             )
+        except pynvml.NVMLError_NotSupported:
+            # This error is expected on some systems.
+            # Only do DEBUG-level logging to reduce noise.
+            logger.debug(
+                f"pyNVML GPU clock speed not supported for device: {device_index}"
+            )
         except Exception:
             logger.exception(
                 f"Failed to get GPU clock speed for device: {device_index}"


Test patch:
diff --git a/tests/unit/utils/test_device_utils.py b/tests/unit/utils/test_device_utils.py
index 941d56f1d8..8e560162e5 100644
--- a/tests/unit/utils/test_device_utils.py
+++ b/tests/unit/utils/test_device_utils.py
@@ -1,7 +1,11 @@
+from unittest.mock import MagicMock, patch
+
 import numpy as np
+import pytest
 import torch
 
 from oumi.utils.device_utils import (
+    _get_nvidia_gpu_runtime_info_impl,
     get_nvidia_gpu_fan_speeds,
     get_nvidia_gpu_memory_utilization,
     get_nvidia_gpu_power_usage,
@@ -213,3 +217,136 @@ def test_nvidia_gpu_runtime_info():
 def test_nvidia_gpu_runtime_info_no_cuda():
     assert get_nvidia_gpu_runtime_info() is None
     log_nvidia_gpu_runtime_info()
+
+
+def _create_mock_pynvml_with_not_supported_error():
+    """Helper function to create a mock pynvml with NVMLError_NotSupported."""
+    mock_pynvml = MagicMock()
+
+    # Mock successful initialization and device count
+    mock_pynvml.nvmlInit.return_value = None
+    mock_pynvml.nvmlDeviceGetCount.return_value = 1
+
+    # Mock successful GPU handle retrieval by default
+    mock_gpu_handle = MagicMock()
+    mock_pynvml.nvmlDeviceGetHandleByIndex.return_value = mock_gpu_handle
+
+    # Create a mock NVMLError_NotSupported exception
+    class MockNVMLError_NotSupported(Exception):
+        pass
+
+    mock_pynvml.NVMLError_NotSupported = MockNVMLError_NotSupported
+
+    return mock_pynvml, MockNVMLError_NotSupported
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        {
+            "name": "utilization",
+            "mock_method": "nvmlDeviceGetUtilizationRates",
+            "function_kwargs": {"utilization": True},
+            "expected_log_message": "GPU utilization not supported for device: 0",
+            "expected_none_fields": ["gpu_utilization", "memory_utilization"],
+            "should_return_none": False,
+        },
+        {
+            "name": "memory",
+            "mock_method": "nvmlDeviceGetMemoryInfo",
+            "function_kwargs": {"memory": True},
+            "expected_log_message": (
+                "pyNVML GPU memory info not supported for device: 0"
+            ),
+            "expected_none_fields": ["used_memory_mb"],
+            "should_return_none": False,
+        },
+        {
+            "name": "temperature",
+            "mock_method": "nvmlDeviceGetTemperature",
+            "function_kwargs": {"temperature": True},
+            "expected_log_message": (
+                "pyNVML GPU temperature not supported for device: 0"
+            ),
+            "expected_none_fields": ["temperature"],
+            "should_return_none": False,
+        },
+        {
+            "name": "fan_speed",
+            "mock_method": "nvmlDeviceGetFanSpeed",
+            "function_kwargs": {"fan_speed": True},
+            "expected_log_message": (
+                "pyNVML GPU fan speed not supported for device: 0"
+            ),
+            "expected_none_fields": ["fan_speed", "fan_speeds"],
+            "should_return_none": False,
+        },
+        {
+            "name": "power_usage",
+            "mock_method": "nvmlDeviceGetPowerUsage",
+            "function_kwargs": {"power_usage": True},
+            "expected_log_message": (
+                "pyNVML GPU power usage not supported for device: 0"
+            ),
+            "expected_none_fields": ["power_usage_watts", "power_limit_watts"],
+            "should_return_none": False,
+        },
+        {
+            "name": "performance_state",
+            "mock_method": "nvmlDeviceGetPerformanceState",
+            "function_kwargs": {"performance_state": True},
+            "expected_log_message": (
+                "pyNVML GPU performance state not supported for device: 0"
+            ),
+            "expected_none_fields": ["performance_state"],
+            "should_return_none": False,
+        },
+        {
+            "name": "clock_speed",
+            "mock_method": "nvmlDeviceGetClockInfo",
+            "function_kwargs": {"clock_speed": True},
+            "expected_log_message": (
+                "pyNVML GPU clock speed not supported for device: 0"
+            ),
+            "expected_none_fields": [
+                "clock_speed_graphics",
+                "clock_speed_sm",
+                "clock_speed_memory",
+            ],
+            "should_return_none": False,
+        },
+    ],
+)
+def test_debug_logging_for_unsupported_gpu_operations(caplog, test_case):
+    """Test that debug logging occurs when GPU operations are not supported."""
+    mock_pynvml, MockNVMLError_NotSupported = (
+        _create_mock_pynvml_with_not_supported_error()
+    )
+
+    # Set up the specific method to raise NotSupported error
+    getattr(
+        mock_pynvml, test_case["mock_method"]
+    ).side_effect = MockNVMLError_NotSupported("Not supported")
+
+    with patch("oumi.utils.device_utils.pynvml", mock_pynvml):
+        with caplog.at_level("DEBUG", logger="oumi"):
+            result = _get_nvidia_gpu_runtime_info_impl(
+                device_index=0, **test_case["function_kwargs"]
+            )
+
+        # Verify debug message was logged
+        assert any(
+            test_case["expected_log_message"] in record.message
+            for record in caplog.records
+            if record.levelname == "DEBUG"
+        ), f"Expected debug message not found: {test_case['expected_log_message']}"
+
+        # Check return value and field values
+        if test_case["should_return_none"]:
+            assert result is None
+        else:
+            assert result is not None
+            for field_name in test_case["expected_none_fields"]:
+                assert getattr(result, field_name) is None, (
+                    f"Expected {field_name} to be None"
+                )