From 3dc64c69281e304c4be341a3724089e501dc8f77 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 29 Sep 2026 17:12:33 +0300 Subject: [PATCH 1/4] fix(engine): size memory pools against the VRAM other processes leave free on Windows --- python/freetoken/engine/graph.py | 18 +++++++- python/freetoken/gpu_select.py | 59 ++++++++++++++++++++----- tests/engine/test_free_memory.py | 75 ++++++++++++++++++++++++++++++++ 3 files changed, 141 insertions(+), 11 deletions(-) create mode 100644 tests/engine/test_free_memory.py diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index eb9684f609..ac278c21f7 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -1,12 +1,14 @@ from __future__ import annotations import gc +import sys from dataclasses import dataclass from typing import TYPE_CHECKING, Dict, List import torch from freetoken.core import Batch, Req, get_global_ctx from freetoken.distributed import get_tp_info +from freetoken.gpu_select import format_gpu_uuid, nvml_free_bytes from freetoken.utils import init_logger, mem_GB from freetoken.utils.progress import emit_progress from tqdm import tqdm @@ -99,7 +101,21 @@ def _determine_cuda_graph_bs( def get_free_memory(device: torch.device) -> int: - return torch.cuda.mem_get_info(device)[0] + """Device memory a pool can take without spilling. + + Under WDDM (Windows) cudaMemGetInfo leaves out what other processes hold, so a card the desktop + uses reads as free and a pool sized to it pages into shared memory; NVML counts every process. + """ + free = torch.cuda.mem_get_info(device)[0] + if sys.platform != "win32": + return free + machine_free = nvml_free_bytes(_device_uuid(device)) + return free if machine_free is None else min(free, machine_free) + + +def _device_uuid(device: torch.device) -> str | None: + index = torch.cuda.current_device() if device.index is None else device.index + return format_gpu_uuid(getattr(torch.cuda.get_device_properties(index), "uuid", None)) class GraphRunner: diff --git a/python/freetoken/gpu_select.py b/python/freetoken/gpu_select.py index 2c7d012b23..a37ed15d67 100644 --- a/python/freetoken/gpu_select.py +++ b/python/freetoken/gpu_select.py @@ -68,12 +68,11 @@ def single_gpu_arg(value: str) -> str: return entries[0] -def _nvml_uuids() -> "list[str] | None": - """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable. +def _load_nvml(): + """The NVML library, or None when no candidate loads. Own ctypes loader instead of torch's _raw_device_uuid_nvml: that helper only knows the Linux library name, raises (not None) when the library is missing, and is private API. NVML exports are cdecl on every platform, so CDLL is right on Windows too (same as nvidia-ml-py). - None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back. """ import ctypes @@ -85,14 +84,24 @@ def _nvml_uuids() -> "list[str] | None": ] else: candidates = ["libnvidia-ml.so.1"] + for name in candidates: + try: + return ctypes.CDLL(name) + except OSError: + continue + return None + + +def _nvml_uuids() -> "list[str] | None": + """Full GPU UUIDs in physical (nvidia-smi) order, or None when NVML is unavailable. + + None on any failure -- no library, a stub library without the _v2 symbols, WSL, a dead device -- and callers fall back. + """ + import ctypes + try: - for name in candidates: - try: - lib = ctypes.CDLL(name) - break - except OSError: - continue - else: + lib = _load_nvml() + if lib is None: return None if lib.nvmlInit() != 0: return None @@ -116,6 +125,36 @@ def _nvml_uuids() -> "list[str] | None": return None +def nvml_free_bytes(uuid: "str | None") -> "int | None": + """Free memory of GPU ``uuid`` across every process, as nvidia-smi reports it; None when NVML cannot say. + + Under WDDM, cudaMemGetInfo leaves out what other processes hold; this does not. + """ + import ctypes + + class Memory(ctypes.Structure): + _fields_ = [("total", ctypes.c_ulonglong), ("free", ctypes.c_ulonglong), ("used", ctypes.c_ulonglong)] + + if uuid is None: + return None + try: + lib = _load_nvml() + if lib is None or lib.nvmlInit() != 0: + return None + try: + handle = ctypes.c_void_p() + if lib.nvmlDeviceGetHandleByUUID(uuid.encode("ascii"), ctypes.byref(handle)) != 0: + return None + memory = Memory() + if lib.nvmlDeviceGetMemoryInfo(handle, ctypes.byref(memory)) != 0: + return None + return int(memory.free) + finally: + lib.nvmlShutdown() + except (OSError, AttributeError): + return None + + def _match_uuid(spec: str, uuids: "list[str]", where: str) -> str: """The unique full UUID that ``spec`` prefixes, else ValueError.""" hits = [u for u in uuids if u.upper().startswith(spec.upper())] diff --git a/tests/engine/test_free_memory.py b/tests/engine/test_free_memory.py new file mode 100644 index 0000000000..eaae672924 --- /dev/null +++ b/tests/engine/test_free_memory.py @@ -0,0 +1,75 @@ +"""The free device memory the engine sizes its pools against. + +Under WDDM (Windows) cudaMemGetInfo leaves out what other processes hold, so a card the desktop uses +reads as free and a pool sized to it pages into shared memory. NVML counts every process. +""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import torch + +from freetoken import gpu_select +from freetoken.engine import graph + +GIB = 1 << 30 +DEVICE_UUID = "80472e39-ad2b-441d-fcff-0f3cc6ab9db2" + + +class FakeNvml: + """NVML with one device and ``free`` bytes free on it; records the UUIDs it is asked about.""" + + def __init__(self, free: int): + self.free = free + self.asked: list[bytes] = [] + + def nvmlInit(self) -> int: + return 0 + + def nvmlShutdown(self) -> int: + return 0 + + def nvmlDeviceGetHandleByUUID(self, uuid: bytes, handle) -> int: + self.asked.append(uuid) + return 0 + + def nvmlDeviceGetMemoryInfo(self, handle, memory) -> int: + memory._obj.free = self.free + return 0 + + +def _machine(monkeypatch, *, cuda_free: int, nvml: FakeNvml | None, platform: str) -> None: + """Fakes CUDA's reading, the device's UUID and the NVML library the engine would load.""" + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (cuda_free, 24 * GIB)) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(uuid=DEVICE_UUID)) + # raising=False: an engine that never reads NVML has no loader, and its free figure shows it + monkeypatch.setattr(gpu_select, "_load_nvml", lambda: nvml, raising=False) + monkeypatch.setattr(sys, "platform", platform) + + +def test_windows_sizes_against_what_other_processes_leave(monkeypatch): + # the reading that spilled: CUDA says 23 GiB free while the desktop holds more than 2 GiB of 24 + nvml = FakeNvml(free=21 * GIB) + _machine(monkeypatch, cuda_free=23 * GIB, nvml=nvml, platform="win32") + assert graph.get_free_memory(torch.device("cuda", 0)) == 21 * GIB + assert nvml.asked == [f"GPU-{DEVICE_UUID}".encode("ascii")] + + +def test_windows_keeps_cudas_figure_when_it_is_the_smaller(monkeypatch): + _machine(monkeypatch, cuda_free=20 * GIB, nvml=FakeNvml(free=21 * GIB), platform="win32") + assert graph.get_free_memory(torch.device("cuda", 0)) == 20 * GIB + + +def test_windows_without_nvml_keeps_cudas_figure(monkeypatch): + _machine(monkeypatch, cuda_free=23 * GIB, nvml=None, platform="win32") + assert graph.get_free_memory(torch.device("cuda", 0)) == 23 * GIB + + +def test_linux_reads_cuda_alone(monkeypatch): + # cudaMemGetInfo already counts every process there, so NVML is not asked + nvml = FakeNvml(free=21 * GIB) + _machine(monkeypatch, cuda_free=23 * GIB, nvml=nvml, platform="linux") + assert graph.get_free_memory(torch.device("cuda", 0)) == 23 * GIB + assert nvml.asked == [] From 04ccee65fb52a963019d0f8cde8197204a58ac4c Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 06:24:57 +0300 Subject: [PATCH 2/4] test(engine): patch the NVML loader strictly _load_nvml exists, so the free-memory tests patch it without raising=False; a renamed loader now fails them at setup. --- tests/engine/test_free_memory.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/engine/test_free_memory.py b/tests/engine/test_free_memory.py index eaae672924..e559d7400f 100644 --- a/tests/engine/test_free_memory.py +++ b/tests/engine/test_free_memory.py @@ -44,8 +44,7 @@ def _machine(monkeypatch, *, cuda_free: int, nvml: FakeNvml | None, platform: st """Fakes CUDA's reading, the device's UUID and the NVML library the engine would load.""" monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (cuda_free, 24 * GIB)) monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(uuid=DEVICE_UUID)) - # raising=False: an engine that never reads NVML has no loader, and its free figure shows it - monkeypatch.setattr(gpu_select, "_load_nvml", lambda: nvml, raising=False) + monkeypatch.setattr(gpu_select, "_load_nvml", lambda: nvml) monkeypatch.setattr(sys, "platform", platform) From f2d6e66dce6d00506f04fef402ed7ec572aa7b58 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 10:31:11 +0300 Subject: [PATCH 3/4] refactor(gpu): derive a visible device's UUID in one place gpu_identity, _visible_of_physical and get_free_memory each spelled the nvidia-smi UUID of a visible device from its torch properties; gpu_uuid is that derivation once. --- python/freetoken/engine/graph.py | 10 +++------- python/freetoken/gpu_select.py | 11 +++++++++-- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index ac278c21f7..6ac612a478 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -8,7 +8,7 @@ import torch from freetoken.core import Batch, Req, get_global_ctx from freetoken.distributed import get_tp_info -from freetoken.gpu_select import format_gpu_uuid, nvml_free_bytes +from freetoken.gpu_select import gpu_uuid, nvml_free_bytes from freetoken.utils import init_logger, mem_GB from freetoken.utils.progress import emit_progress from tqdm import tqdm @@ -109,13 +109,9 @@ def get_free_memory(device: torch.device) -> int: free = torch.cuda.mem_get_info(device)[0] if sys.platform != "win32": return free - machine_free = nvml_free_bytes(_device_uuid(device)) - return free if machine_free is None else min(free, machine_free) - - -def _device_uuid(device: torch.device) -> str | None: index = torch.cuda.current_device() if device.index is None else device.index - return format_gpu_uuid(getattr(torch.cuda.get_device_properties(index), "uuid", None)) + machine_free = nvml_free_bytes(gpu_uuid(index)) + return free if machine_free is None else min(free, machine_free) class GraphRunner: diff --git a/python/freetoken/gpu_select.py b/python/freetoken/gpu_select.py index a37ed15d67..a06fc32577 100644 --- a/python/freetoken/gpu_select.py +++ b/python/freetoken/gpu_select.py @@ -262,7 +262,7 @@ def _visible_of_physical(uuid: str) -> int: seen: list[str] = [] hits: list[int] = [] for v in range(torch.cuda.device_count()): - u = format_gpu_uuid(getattr(torch.cuda.get_device_properties(v), "uuid", None)) + u = gpu_uuid(v) seen.append(u or "?") if u is not None and u.upper().startswith(uuid.upper()): hits.append(v) @@ -313,6 +313,13 @@ def format_gpu_uuid(raw) -> str | None: return None if raw is None else f"{UUID_PREFIX}{raw}" +def gpu_uuid(index: int) -> str | None: + """nvidia-smi form GPU- of visible device ``index``; None where torch reports no uuid.""" + import torch + + return format_gpu_uuid(getattr(torch.cuda.get_device_properties(index), "uuid", None)) + + def gpu_identity(index: int) -> dict: """{index, name, uuid, total_bytes} of visible device ``index``.""" import torch @@ -321,6 +328,6 @@ def gpu_identity(index: int) -> dict: return { "index": index, "name": props.name, - "uuid": format_gpu_uuid(getattr(props, "uuid", None)), + "uuid": gpu_uuid(index), "total_bytes": int(props.total_memory), } From 35088b03b46d8d7028c4537b5c4ed01321d8f84c Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 12:12:28 +0300 Subject: [PATCH 4/4] docs(gpu): name the NVML free-memory read in gpu_select, and the platform the test fakes --- python/freetoken/gpu_select.py | 1 + tests/engine/test_free_memory.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/python/freetoken/gpu_select.py b/python/freetoken/gpu_select.py index a06fc32577..9b45c709cb 100644 --- a/python/freetoken/gpu_select.py +++ b/python/freetoken/gpu_select.py @@ -9,6 +9,7 @@ Each worker publishes its own entry (set_assigned_gpu / assign_gpu) and binds it when CUDA comes up (bind_assigned_gpu) by matching the UUID against CUDA's visible devices. One process runs on one GPU. Binding is unconditional: a process that publishes nothing binds a default ordinal and records it, so assigned_visible_gpu() names that card in every case. No process mutates CUDA_VISIBLE_DEVICES, and the UUID match holds under any CUDA_DEVICE_ORDER. +NVML also reports a GPU's free memory across every process (nvml_free_bytes), which WDDM's cudaMemGetInfo does not. Stdlib only (torch is imported lazily); not under freetoken.utils, which imports transformers. """ diff --git a/tests/engine/test_free_memory.py b/tests/engine/test_free_memory.py index e559d7400f..3b3a4f6f0e 100644 --- a/tests/engine/test_free_memory.py +++ b/tests/engine/test_free_memory.py @@ -41,7 +41,7 @@ def nvmlDeviceGetMemoryInfo(self, handle, memory) -> int: def _machine(monkeypatch, *, cuda_free: int, nvml: FakeNvml | None, platform: str) -> None: - """Fakes CUDA's reading, the device's UUID and the NVML library the engine would load.""" + """Fakes CUDA's reading, the device's UUID, the NVML library the engine would load and the platform.""" monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (cuda_free, 24 * GIB)) monkeypatch.setattr(torch.cuda, "get_device_properties", lambda index: SimpleNamespace(uuid=DEVICE_UUID)) monkeypatch.setattr(gpu_select, "_load_nvml", lambda: nvml)