Normalize GPU VRAM against registered capacities

src/gpuhunt/_internal/constraints.py

@@ -1,3 +1,4 @@
+import logging
 import re
 from collections.abc import Container, Iterable
 from typing import TypeVar
@@ -15,6 +16,10 @@
     TPUInfo,
 )
 
+logger = logging.getLogger(__name__)
+
+GPU_MEMORY_TOLERANCE = 0.07
+
 
 def matches(i: CatalogItem, q: QueryFilter) -> bool:
     """
@@ -108,25 +113,54 @@ def get_gpu_vendor(gpu_name: str | None) -> AcceleratorVendor | None:
     return None
 
 
-def correct_gpu_memory_gib(gpu_name: str, memory_mib: float) -> int:
-    """
-    Round to whole number of gibibytes and attempt correcting the reported GPU
-    memory size if the actual memory size for that GPU is known and the
-    difference between the reported and the known memory is within a heuristic
-    threshold.
+def correct_gpu_memory_gib(gpu_name: str, memory_mib: float, *, context: str | None = None) -> int:
+    """Normalize full-GPU, per-device MiB to whole GiB using KNOWN_ACCELERATORS.
 
-    This is useful for cases when nvidia-smi or cloud providers report the GPU
-    memory imprecisely.
+    Preserve exact capacities. Otherwise, return a registered capacity only if the
+    reported value is within ±7% of exactly one registered capacity for gpu_name.
+    Unknown models and unmatched or ambiguous values warn and are rounded to whole GiB.
+    Callers must exclude partial GPU allocations. context identifies the provider and
+    offer (or host/node) in warnings only.
     """
-
-    memory_gib = memory_mib / 1024
-    known_memories_gib = {gpu.memory for gpu in find_accelerators(names=[gpu_name])}
-    if known_memories_gib:
-        closest_known_memory_gib = min(known_memories_gib, key=lambda x: abs(x - memory_gib))
-        difference_gib = abs(closest_known_memory_gib - memory_gib)
-        if difference_gib / closest_known_memory_gib < 0.07:
-            return closest_known_memory_gib
-    return round(memory_gib)
+    reported_memory_gib = memory_mib / 1024
+    rounded_memory_gib = round(reported_memory_gib)
+    context = f" ({context})" if context else ""
+    memories = {gpu.memory for gpu in find_accelerators(names=[gpu_name]) if gpu.memory > 0}
+    if not memories:
+        logger.warning(
+            "Unknown GPU %s%s: model is not in KNOWN_ACCELERATORS; "
+            "reported VRAM=%s GiB per GPU; using rounded VRAM=%s GiB per GPU",
+            gpu_name,
+            context,
+            reported_memory_gib,
+            rounded_memory_gib,
+        )
+        return rounded_memory_gib
+    if reported_memory_gib in memories:
+        return rounded_memory_gib
+    candidates = [
+        memory
+        for memory in memories
+        if memory * (1 - GPU_MEMORY_TOLERANCE)
+        <= reported_memory_gib
+        <= memory * (1 + GPU_MEMORY_TOLERANCE)
+    ]
+    if len(candidates) == 1:
+        return candidates[0]
+    logger.warning(
+        "Cannot normalize VRAM for GPU %s%s: reported VRAM=%s GiB per GPU; "
+        "%s within ±%g%%; registered capacities=%s GiB; using rounded VRAM=%s GiB per GPU",
+        gpu_name,
+        context,
+        reported_memory_gib,
+        "no registered capacity matches"
+        if not candidates
+        else "multiple registered capacities match",
+        GPU_MEMORY_TOLERANCE * 100,
+        sorted(memories),
+        rounded_memory_gib,
+    )
+    return rounded_memory_gib
 
 
 def is_nvidia_superchip(gpu_name: str) -> bool:
@@ -146,6 +180,8 @@ def is_nvidia_superchip(gpu_name: str) -> bool:
     NvidiaGPUInfo(name="A100", memory=40, compute_capability=(8, 0)),
     NvidiaGPUInfo(name="A100", memory=80, compute_capability=(8, 0)),
     NvidiaGPUInfo(name="A10G", memory=24, compute_capability=(8, 6)),
+    NvidiaGPUInfo(name="A2000", memory=6, compute_capability=(8, 6)),
+    NvidiaGPUInfo(name="A2000", memory=12, compute_capability=(8, 6)),
     NvidiaGPUInfo(name="A4000", memory=16, compute_capability=(8, 6)),
     NvidiaGPUInfo(name="A4500", memory=20, compute_capability=(8, 6)),
     NvidiaGPUInfo(name="A5000", memory=24, compute_capability=(8, 6)),
@@ -199,10 +235,16 @@ def is_nvidia_superchip(gpu_name: str) -> bool:
     NvidiaGPUInfo(name="RTX5080", memory=16, compute_capability=(12, 0)),
     NvidiaGPUInfo(name="RTX5090", memory=32, compute_capability=(12, 0)),
     NvidiaGPUInfo(name="RTX6000", memory=24, compute_capability=(7, 5)),
+    NvidiaGPUInfo(name="RTX8000", memory=48, compute_capability=(7, 5)),
     NvidiaGPUInfo(name="RTX2000Ada", memory=16, compute_capability=(8, 9)),
     NvidiaGPUInfo(name="RTX4000Ada", memory=20, compute_capability=(8, 9)),
+    NvidiaGPUInfo(name="RTX4500Ada", memory=24, compute_capability=(8, 9)),
+    NvidiaGPUInfo(name="RTX5000Ada", memory=32, compute_capability=(8, 9)),
     NvidiaGPUInfo(name="RTX6000Ada", memory=48, compute_capability=(8, 9)),
+    NvidiaGPUInfo(name="RTXPRO4000", memory=24, compute_capability=(12, 0)),
     NvidiaGPUInfo(name="RTXPRO4500", memory=32, compute_capability=(12, 0)),
+    NvidiaGPUInfo(name="RTXPRO5000", memory=48, compute_capability=(12, 0)),
+    NvidiaGPUInfo(name="RTXPRO5000", memory=72, compute_capability=(12, 0)),
     NvidiaGPUInfo(name="RTXPRO6000", memory=96, compute_capability=(12, 0)),
     NvidiaGPUInfo(name="T4", memory=16, compute_capability=(7, 5)),
     NvidiaGPUInfo(name="V100", memory=16, compute_capability=(7, 0)),
@@ -259,7 +301,7 @@ def is_nvidia_superchip(gpu_name: str) -> bool:
     ),
     AMDGPUInfo(
         name="MI325X",
-        memory=288,
+        memory=256,
         architecture=AMDArchitecture.CDNA3,
         device_ids=(0x74A5,),
     ),

src/gpuhunt/providers/aws.py

@@ -17,6 +17,7 @@
 from botocore.config import Config
 from botocore.exceptions import ClientError, ConnectTimeoutError, EndpointConnectionError
 
+from gpuhunt._internal.constraints import correct_gpu_memory_gib
 from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter
 from gpuhunt.providers.base import OfflineProvider
 
@@ -223,7 +224,10 @@ def _get_gpu_details(self, rows: list["_PricingRow"]) -> dict[str, tuple[str, fl
                                 gpus[i["InstanceType"]] = (
                                     GPU_NAME_MAPPING.get(gpu["Name"]) or gpu["Name"],
                                     _get_gpu_memory_gib(
-                                        gpu["Name"], gpu["MemoryInfo"]["SizeInMiB"]
+                                        gpu["Name"],
+                                        gpu["MemoryInfo"]["SizeInMiB"],
+                                        context=f"provider={self.NAME}, instance={i['InstanceType']}, "
+                                        f"location={region}",
                                     ),
                                 )
             except ConnectTimeoutError as e:
@@ -468,20 +472,12 @@ def _make_offer(row: _PricingRow, gpus: dict[str, tuple[str, float]]) -> Catalog
     )
 
 
-def _get_gpu_memory_gib(gpu_name: str, reported_memory_mib: int) -> float:
-    """
-    Fixes L4 memory size misreported by AWS API
-    """
-
-    if gpu_name != "L4":
-        return reported_memory_mib / 1024
-
-    if reported_memory_mib not in (22888, 91553, 183105):
-        logger.warning(
-            "The L4 memory size reported by AWS changed. "
-            "Please check that it is now correct and remove the hardcoded size if it is."
-        )
-    return 24
+def _get_gpu_memory_gib(
+    gpu_name: str, reported_memory_mib: int, *, context: str | None = None
+) -> int:
+    """Resolve the EC2 model name and normalize its per-device MiB to whole GiB."""
+    gpu_name = GPU_NAME_MAPPING.get(gpu_name, gpu_name)
+    return correct_gpu_memory_gib(gpu_name, reported_memory_mib, context=context)
 
 
 def _prefilter_rows(lines: Iterable[str]) -> Iterator[str]:

src/gpuhunt/providers/digitalocean.py

@@ -3,7 +3,7 @@
 
 import requests
 
-from gpuhunt._internal.constraints import get_gpu_vendor
+from gpuhunt._internal.constraints import correct_gpu_memory_gib, get_gpu_vendor
 from gpuhunt._internal.models import CatalogItem, QueryFilter
 from gpuhunt.providers.base import OnlineProvider, get_creds_env
 
@@ -64,7 +64,8 @@ def _make_offers(response) -> list[CatalogItem]:
         if gpu_info:
             gpu_count = gpu_info["count"]
             gpu_vram_info = gpu_info["vram"]
-            gpu_memory = float(gpu_vram_info["amount"])
+            # vram.amount is the combined VRAM of all GPUs, in GiB.
+            gpu_memory = float(gpu_vram_info["amount"]) / gpu_count
             gpu_model = gpu_info["model"]
             # gpu_model uses patterns like "amd_mi300x", "nvidia_h100", "nvidia_rtx6000_ada"
             model_parts = gpu_model.split("_")
@@ -75,6 +76,11 @@ def _make_offers(response) -> list[CatalogItem]:
                     f"Could not determine GPU vendor for model '{gpu_model}'. Skipping droplet '{size['slug']}'."
                 )
                 continue
+            gpu_memory = correct_gpu_memory_gib(
+                gpu_name,
+                gpu_memory * 1024,
+                context=f"provider={DigitalOceanProvider.NAME}, instance={size['slug']}",
+            )
         else:
             gpu_count = 0
             gpu_vendor = None

src/gpuhunt/providers/lium.py

@@ -114,11 +114,13 @@ def _make_offers(node: dict) -> list[CatalogItem]:
         return []
     gpu_name = get_dstack_gpu_name(gpu_model)
     if not find_accelerators(names=[gpu_name], vendors=[AcceleratorVendor.NVIDIA]):
-        # Like the seeweb and jarvislabs providers, keep unmapped models out of the catalog
-        # rather than publishing a name nothing can match; the warning makes them easy to spot.
         logger.warning("Skipping Lium node %s: unknown gpu_model %r", node_id, gpu_model)
         return []
-    gpu_memory = correct_gpu_memory_gib(gpu_name, gpu_memory_gb * 1024)
+    gpu_memory = correct_gpu_memory_gib(
+        gpu_name,
+        gpu_memory_gb * 1024,
+        context=f"provider={LiumProvider.NAME}, node={node_id}, location={location}",
+    )
     provider_data = LiumCatalogItemProviderData()
     rent_url = node.get("rent_url")
     if rent_url:

src/gpuhunt/providers/oci.py

@@ -12,6 +12,7 @@
 from typing_extensions import TypedDict
 
 from gpuhunt._internal.constraints import (
+    correct_gpu_memory_gib,
     find_accelerators,
     get_gpu_vendor,
     is_nvidia_superchip,
@@ -309,9 +310,14 @@ def shape_to_resources(
             vcpus = qty if cpu_arch == CPUArchitecture.ARM else qty * 2
             if shape.gpu_qty:
                 # For GPU shapes the "ocpu" product is priced per GPU-hour
+                gpu_memory = shape.get_gpu_unit_memory_gb()
+                if gpu_name is not None and gpu_memory is not None:
+                    gpu_memory = correct_gpu_memory_gib(
+                        gpu_name, gpu_memory * 1024, context=f"provider=oci, shape={shape.name}"
+                    )
                 gpu = GPUConfiguration(
                     units_count=shape.gpu_qty,
-                    unit_memory_gb=shape.get_gpu_unit_memory_gb(),
+                    unit_memory_gb=gpu_memory,
                     name=gpu_name,
                     vendor=get_gpu_vendor(gpu_name),
                     price=product_price * shape.gpu_qty,

src/gpuhunt/providers/runpod.py

@@ -6,7 +6,7 @@
 from requests import RequestException
 from typing_extensions import NotRequired, TypedDict
 
-from gpuhunt._internal.constraints import find_accelerators
+from gpuhunt._internal.constraints import correct_gpu_memory_gib, find_accelerators
 from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter
 from gpuhunt.providers.base import OfflineProvider
 
@@ -164,7 +164,11 @@ def _make_offers(self, query_variables: dict, pod: dict) -> list[CatalogItem]:
                 gpu_vendor=listed_gpu_vendor_and_name[0],
                 gpu_count=lowest_price_input_variables["gpuCount"],
                 gpu_name=listed_gpu_vendor_and_name[1],
-                gpu_memory=pod["memoryInGb"],
+                gpu_memory=correct_gpu_memory_gib(
+                    listed_gpu_vendor_and_name[1],
+                    pod["memoryInGb"] * 1024,
+                    context=f"provider={self.NAME}, instance={pod['id']}, location={location}",
+                ),
                 spot=False,
                 disk_size=None,
                 provider_data={},
@@ -211,7 +215,12 @@ def _fetch_cluster_offers(self) -> list[CatalogItem]:
                     gpu_vendor=gpu_vendor,
                     gpu_count=pod_type["maxGpuCountSecureCloud"],
                     gpu_name=gpu_name,
-                    gpu_memory=pod_type["memoryInGb"],
+                    gpu_memory=correct_gpu_memory_gib(
+                        gpu_name,
+                        pod_type["memoryInGb"] * 1024,
+                        context=f"provider={self.NAME}, instance={pod_type['id']}, "
+                        f"location={location['id']}, cluster=True",
+                    ),
                     spot=False,
                     disk_size=None,
                     flags=["runpod-cluster"],

src/gpuhunt/providers/seeweb.py

@@ -3,7 +3,7 @@
 
 import requests
 
-from gpuhunt._internal.constraints import find_accelerators
+from gpuhunt._internal.constraints import correct_gpu_memory_gib, find_accelerators
 from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter
 from gpuhunt.providers.base import OnlineProvider, get_creds_env
 
@@ -101,6 +101,9 @@ def _make_offers(plan: dict) -> list[CatalogItem]:
         logger.warning("Skipping Seeweb plan %s: unknown gpu_label %r", plan_name, gpu_label)
         return []
     gpu_name, gpu_memory = resolved
+    gpu_memory = correct_gpu_memory_gib(
+        gpu_name, gpu_memory * 1024, context=f"provider={SeewebProvider.NAME}, plan={plan_name}"
+    )
 
     try:
         cpu = _parse_int(plan.get("cpu"))

src/gpuhunt/providers/vastai.py

@@ -8,7 +8,7 @@
 import requests
 from typing_extensions import NotRequired, TypedDict
 
-from gpuhunt._internal.constraints import correct_gpu_memory_gib
+from gpuhunt._internal.constraints import GPU_MEMORY_TOLERANCE, correct_gpu_memory_gib
 from gpuhunt._internal.models import AcceleratorVendor, CatalogItem, QueryFilter
 from gpuhunt.providers.base import OnlineProvider
 
@@ -76,12 +76,18 @@ def get(
                 logger.warning("Offer %s does not satisfy filters", offer["id"])
                 continue
             gpu_name = get_dstack_gpu_name(offer["gpu_name"])
-            gpu_memory = correct_gpu_memory_gib(gpu_name, offer["gpu_ram"])
+            location = get_location(offer["geolocation"])
+            # Vast reports measured VRAM in MiB.
+            gpu_memory = correct_gpu_memory_gib(
+                gpu_name,
+                offer["gpu_ram"],
+                context=f"provider={self.NAME}, offer={offer['id']}, location={location}",
+            )
             disk_cost = disk_size * offer["storage_cost"] / 30 / 24
             ondemand_offer = CatalogItem(
                 provider=VastAIProvider.NAME,
                 instance_name=str(offer["id"]),
-                location=get_location(offer["geolocation"]),
+                location=location,
                 # storage_cost is $/gb/month
                 price=round(offer["dph_base"] + disk_cost, 5),
                 cpu=cpu,
@@ -136,11 +142,10 @@ def make_filters(self, q: QueryFilter) -> dict[str, Any]:
             else:
                 # If GPU name mapping fails, fetch all offers (to filter locally)
                 filters["limit"] = 3000
-        # See correct_gpu_memory_gib in gpuhunt/_internal/constraints.py
         if q.min_gpu_memory is not None:
-            filters["gpu_ram"]["gte"] = q.min_gpu_memory * 1024 * 0.93
+            filters["gpu_ram"]["gte"] = q.min_gpu_memory * (1 - GPU_MEMORY_TOLERANCE) * 1024
         if q.max_gpu_memory is not None:
-            filters["gpu_ram"]["lte"] = q.max_gpu_memory * 1024 * 1.07
+            filters["gpu_ram"]["lte"] = q.max_gpu_memory * (1 + GPU_MEMORY_TOLERANCE) * 1024
         if q.min_disk_size is not None:
             filters["disk_space"]["gte"] = q.min_disk_size
         if q.max_disk_size is not None:
@@ -213,6 +218,17 @@ class VastAICatalogItemProviderData(TypedDict):
     "H100": ["H100 PCIE", "H100 SXM"],
     "H100NVL": ["H100 NVL"],
     "B300": ["B300", "B300 PC"],
+    "RTX2060SUPER": ["RTX 2060S"],
+    "RTX2070SUPER": ["RTX 2070S"],
+    "RTX4070SUPER": ["RTX 4070S"],
+    "RTX4070TiSUPER": ["RTX 4070S Ti"],
+    "RTX4080SUPER": ["RTX 4080S"],
+    "RTX4000": ["Q RTX 4000", "RTX 4000"],
+    "RTX5000": ["Q RTX 5000", "RTX 5000"],
+    "RTX6000": ["Q RTX 6000", "RTX 6000"],
+    "RTX8000": ["Q RTX 8000", "RTX 8000"],
+    "RTXPRO4000": ["RTX PRO 4000"],
+    "RTXPRO5000": ["RTX PRO 5000"],
 }
 
 GPU_MAPPING_RULES = {
@@ -234,18 +250,15 @@ def get_vastai_gpu_names(gpu_name: str) -> list[str]:
 
 
 def get_dstack_gpu_name(gpu_name: str) -> str:
-    """
-    Convert VastAI GPU names to a standardized format using essential heuristics
-    """
+    """Convert a Vast product label to the catalog model name."""
+    for name, labels in GPU_MAPPING.items():
+        if gpu_name in labels:
+            return name
     gpu_name = gpu_name.replace("RTX A", "A").replace("Tesla ", "")
     if gpu_name.startswith("A100 "):
         return "A100"
     if gpu_name.startswith("H100 ") and "NVL" not in gpu_name:
         return "H100"
-    # "B300 PC" is how Vast lists boards that report "NVIDIA B300 SXM6 PC" rather than
-    # "NVIDIA B300 SXM6 AC": the same B300 GPU (PCI device 0x3182) on a different board SKU
-    if gpu_name == "B300 PC":
-        return "B300"
     return gpu_name.replace(" ", "")