From d30bd40ae78d459073b2b5856267b7a39adb1746 Mon Sep 17 00:00:00 2001 From: mschwab Date: Wed, 22 Jul 2026 22:50:00 -0700 Subject: [PATCH 1/2] fix(llm): size VRAM headroom from reclaimable memory on integrated GPUs `_get_vram_allocations` derived available VRAM from `torch.cuda.mem_get_info().free`, which on integrated GPUs (e.g. NVIDIA GB10 / Grace, where GPU memory is system memory) counts only unallocated pages and ignores the reclaimable page cache. On such a device with most memory in buff/cache, `free` reads a couple GiB; after the 2 GiB safety buffer the usable fraction collapses to 0, and the `gpu.vram` preflight check hard-fails every job with "exceeds available ~0.0 GiB" even when tens of GiB are actually reclaimable. On integrated GPUs, use the kernel's `MemAvailable` (which accounts for reclaimable memory), capped at device total, instead of `mem_get_info`'s free figure. Discrete-GPU behavior is unchanged. Falls back to the CUDA free value when `/proc/meminfo` is unavailable (e.g. non-Linux hosts). Signed-off-by: mschwab --- src/nemo_safe_synthesizer/llm/utils.py | 27 +++++++++++++++ tests/llm/test_utils.py | 46 +++++++++++++++++++++++++- 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/src/nemo_safe_synthesizer/llm/utils.py b/src/nemo_safe_synthesizer/llm/utils.py index 5ac6ed0ba..1c9cc31b5 100644 --- a/src/nemo_safe_synthesizer/llm/utils.py +++ b/src/nemo_safe_synthesizer/llm/utils.py @@ -436,12 +436,35 @@ class _VRAMAllocation: memory_bytes: int +def _reclaimable_available_bytes() -> int | None: + """Return kernel-reclaimable system memory (``MemAvailable``) in bytes. + + ``MemAvailable`` in ``/proc/meminfo`` accounts for reclaimable page cache, so + it reflects what the kernel can actually hand out. Returns ``None`` when it + cannot be read (for example on non-Linux hosts, where ``/proc/meminfo`` is + absent). + """ + try: + with open("/proc/meminfo", encoding="utf-8") as meminfo: + for line in meminfo: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) * 1024 # reported in kibibytes + except (OSError, ValueError, IndexError): + return None + return None + + def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _VRAMAllocation]: """Calculate maximum memory allocation for each available GPU. Reserves a 2 GiB safety buffer on each device, then applies ``max_vram_fraction`` to the remaining free memory. + On systems with integrated GPUs (e.g. DGX Spark) GPU memory *is* system memory, + so ``torch.cuda.mem_get_info`` reports only unallocated pages and ignores the + reclaimable page cache -- badly under-reporting real headroom. There we use + the kernel's ``MemAvailable`` (capped at device total) instead. + Args: max_vram_fraction: Fraction of total GPU memory to allocate. Defaults to ``0.8`` (80 %). @@ -459,6 +482,10 @@ def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _ num_gpus = torch.cuda.device_count() for i in range(num_gpus): free, total = torch.cuda.mem_get_info(device=i) + if getattr(torch.cuda.get_device_properties(i), "is_integrated", False): + available = _reclaimable_available_bytes() + if available is not None: + free = min(max(free, available), total) safe_free = max(free - (2 * 1024**3), 0) gpu_memory_utilization = min(max_vram_fraction, safe_free / total) if total > 0 else 0.0 memory_bytes = int(gpu_memory_utilization * total) diff --git a/tests/llm/test_utils.py b/tests/llm/test_utils.py index 05c6b3b90..4e48521eb 100644 --- a/tests/llm/test_utils.py +++ b/tests/llm/test_utils.py @@ -4,12 +4,13 @@ """Unit tests for llm.utils helpers.""" from pathlib import Path -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, mock_open, patch import pytest from nemo_safe_synthesizer.llm.utils import ( ModelRef, + _reclaimable_available_bytes, get_max_memory_map, get_max_vram, get_quantization_config, @@ -73,15 +74,58 @@ def test_trust_remote_code_for_model_requires_configured_cache_root(tmp_path: Pa def test_vram_helpers_return_fraction_and_hf_memory_map() -> None: gib = 1024**3 + discrete_gpu = MagicMock(is_integrated=False) with ( patch("torch.cuda.is_available", return_value=True), patch("torch.cuda.device_count", return_value=1), patch("torch.cuda.mem_get_info", return_value=(10 * gib, 16 * gib)), + patch("torch.cuda.get_device_properties", return_value=discrete_gpu), ): assert get_max_vram(max_vram_fraction=0.8) == {0: 0.5} assert get_max_memory_map(max_vram_fraction=0.8) == {0: 8 * gib} +def test_vram_helpers_use_reclaimable_memory_on_integrated_gpus() -> None: + gib = 1024**3 + integrated_gpu = MagicMock(is_integrated=True) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.device_count", return_value=1), + # Unified memory: mem_get_info under-reports free (2 GiB) ... + patch("torch.cuda.mem_get_info", return_value=(2 * gib, 120 * gib)), + patch("torch.cuda.get_device_properties", return_value=integrated_gpu), + # ... but 60 GiB is reclaimable/available per the kernel. + patch("nemo_safe_synthesizer.llm.utils._reclaimable_available_bytes", return_value=60 * gib), + ): + # free -> 60 GiB, safe_free -> 58 GiB, utilization -> 58/120. + assert get_max_vram(max_vram_fraction=0.8) == {0: pytest.approx(58 / 120)} + assert get_max_memory_map(max_vram_fraction=0.8) == {0: 58 * gib} + + +def test_vram_helpers_fall_back_to_cuda_free_when_meminfo_unreadable() -> None: + gib = 1024**3 + integrated_gpu = MagicMock(is_integrated=True) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.device_count", return_value=1), + patch("torch.cuda.mem_get_info", return_value=(10 * gib, 16 * gib)), + patch("torch.cuda.get_device_properties", return_value=integrated_gpu), + patch("nemo_safe_synthesizer.llm.utils._reclaimable_available_bytes", return_value=None), + ): + assert get_max_vram(max_vram_fraction=0.8) == {0: 0.5} + + +def test_reclaimable_available_bytes_parses_meminfo() -> None: + meminfo = "MemTotal: 127603160 kB\nMemFree: 204 kB\nMemAvailable: 68157440 kB\n" + with patch("builtins.open", mock_open(read_data=meminfo)): + assert _reclaimable_available_bytes() == 68157440 * 1024 + + +def test_reclaimable_available_bytes_returns_none_when_absent() -> None: + with patch("builtins.open", side_effect=FileNotFoundError): + assert _reclaimable_available_bytes() is None + + def test_model_ref_trusts_snapshot_under_configured_cache_root(tmp_path: Path, hf_cached_snapshot_factory) -> None: """HF cache path recognition intentionally tracks Hub snapshot metadata.""" cache_root = tmp_path / "custom-cache" From 9df2ba92f3a0cca0a737da6b63c6503ae5a53ab8 Mon Sep 17 00:00:00 2001 From: mschwab Date: Wed, 22 Jul 2026 22:50:00 -0700 Subject: [PATCH 2/2] fix(llm): hoist MemAvailable read and test the device-total cap Address PR review: read `MemAvailable` once above the per-GPU loop (it is a system-wide value) rather than per device, and add a regression test for the `min(free, total)` cap where `MemAvailable` exceeds device total, guarding the 2 GiB safety buffer. Signed-off-by: mschwab --- src/nemo_safe_synthesizer/llm/utils.py | 8 ++++---- tests/llm/test_utils.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/src/nemo_safe_synthesizer/llm/utils.py b/src/nemo_safe_synthesizer/llm/utils.py index 1c9cc31b5..f3b6ead41 100644 --- a/src/nemo_safe_synthesizer/llm/utils.py +++ b/src/nemo_safe_synthesizer/llm/utils.py @@ -480,12 +480,12 @@ def _get_vram_allocations(max_vram_fraction: float | None = None) -> dict[int, _ if torch.cuda.is_available(): num_gpus = torch.cuda.device_count() + # System-wide value; read once rather than per device. + reclaimable = _reclaimable_available_bytes() for i in range(num_gpus): free, total = torch.cuda.mem_get_info(device=i) - if getattr(torch.cuda.get_device_properties(i), "is_integrated", False): - available = _reclaimable_available_bytes() - if available is not None: - free = min(max(free, available), total) + if getattr(torch.cuda.get_device_properties(i), "is_integrated", False) and reclaimable is not None: + free = min(max(free, reclaimable), total) safe_free = max(free - (2 * 1024**3), 0) gpu_memory_utilization = min(max_vram_fraction, safe_free / total) if total > 0 else 0.0 memory_bytes = int(gpu_memory_utilization * total) diff --git a/tests/llm/test_utils.py b/tests/llm/test_utils.py index 4e48521eb..627004f46 100644 --- a/tests/llm/test_utils.py +++ b/tests/llm/test_utils.py @@ -102,6 +102,22 @@ def test_vram_helpers_use_reclaimable_memory_on_integrated_gpus() -> None: assert get_max_memory_map(max_vram_fraction=0.8) == {0: 58 * gib} +def test_vram_helpers_cap_reclaimable_memory_at_device_total() -> None: + gib = 1024**3 + integrated_gpu = MagicMock(is_integrated=True) + with ( + patch("torch.cuda.is_available", return_value=True), + patch("torch.cuda.device_count", return_value=1), + patch("torch.cuda.mem_get_info", return_value=(2 * gib, 120 * gib)), + patch("torch.cuda.get_device_properties", return_value=integrated_gpu), + # MemAvailable exceeds device total; must be capped so the 2 GiB buffer survives. + patch("nemo_safe_synthesizer.llm.utils._reclaimable_available_bytes", return_value=200 * gib), + ): + # free -> 120 GiB (capped), safe_free -> 118 GiB, utilization -> 118/120. + assert get_max_vram(max_vram_fraction=1.0) == {0: pytest.approx(118 / 120)} + assert get_max_memory_map(max_vram_fraction=1.0) == {0: 118 * gib} + + def test_vram_helpers_fall_back_to_cuda_free_when_meminfo_unreadable() -> None: gib = 1024**3 integrated_gpu = MagicMock(is_integrated=True)