OutOfMemoryError in single L40 GPU

[dianjingshilezhi]

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 40.42 GiB. GPU 0 has a total capacity of 44.40 GiB of which 33.26 GiB is free. Including non-PyTorch memory, this process has 11.13 GiB memory in use. Of the allocated memory 10.25 GiB is allocated by PyTorch, and 394.50 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

accelerate 1.13.0
certifi 2026.5.20
charset-normalizer 3.4.7
cuda-bindings 13.3.1
cuda-pathfinder 1.5.5
cuda-toolkit 12.8.0
decord 0.6.0
filelock 3.29.1
fsspec 2026.4.0
hf-xet 1.5.0
huggingface_hub 0.36.2
idna 3.18
Jinja2 3.1.6
lmdb 1.7.5
MarkupSafe 3.0.3
mpmath 1.3.0
networkx 3.6.1
numpy 1.25.0
nvidia-cublas 13.1.0.3
nvidia-cublas-cu12 12.8.4.1
nvidia-cuda-cupti 13.0.85
nvidia-cuda-cupti-cu12 12.8.90
nvidia-cuda-nvrtc 13.0.88
nvidia-cuda-nvrtc-cu12 12.8.93
nvidia-cuda-runtime 13.0.96
nvidia-cuda-runtime-cu12 12.8.90
nvidia-cudnn-cu12 9.10.2.21
nvidia-cudnn-cu13 9.19.0.56
nvidia-cufft 12.0.0.61
nvidia-cufft-cu12 11.3.3.83
nvidia-cufile 1.15.1.6
nvidia-cufile-cu12 1.13.1.3
nvidia-curand 10.4.0.35
nvidia-curand-cu12 10.3.9.90
nvidia-cusolver 12.0.4.66
nvidia-cusolver-cu12 11.7.3.90
nvidia-cusparse 12.6.3.3
nvidia-cusparse-cu12 12.5.8.93
nvidia-cusparselt-cu12 0.7.1
nvidia-cusparselt-cu13 0.8.0
nvidia-nccl-cu12 2.27.3
nvidia-nccl-cu13 2.28.9
nvidia-nvjitlink 13.0.88
nvidia-nvjitlink-cu12 12.8.93
nvidia-nvshmem-cu13 3.4.5
nvidia-nvtx 13.0.85
nvidia-nvtx-cu12 12.8.90
opencv-python-headless 4.11.0.86
packaging 26.2
peft 0.19.1
pillow 11.1.0
pip 25.2
psutil 7.2.2
PyYAML 6.0.3
regex 2026.5.9
requests 2.34.2
safetensors 0.7.0
setuptools 80.9.0
sympy 1.14.0
tokenizers 0.22.2
torch 2.8.0
torchaudio 2.8.0
torchvision 0.23.0
tqdm 4.67.3
transformers 4.57.1
triton 3.4.0
typing_extensions 4.15.0
urllib3 2.7.0
wheel 0.45.1

[michaelmanly]

That looks like a hard VRAM ceiling rather than a normal code bug.

What exact command / batch size / input resolution are you running when the L40 OOMs?

If you can share the run command, it may be easier to test the same workload on a larger-memory GPU or with a capped run so it fails safely instead of wasting the whole session. @dianjingshilezhi

[dianjingshilezhi]

I use the sample script to infer the model.

import re
import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer, AutoProcessor

class LocateAnythingWorker:
    """Stateful worker that loads the model once and serves perception queries."""
    def __init__(self, model_path: str, device: str = "cuda", dtype=torch.bfloat16):
        self.device = device
        self.dtype = dtype
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        self.processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
        self.model = AutoModel.from_pretrained(
            model_path,
            torch_dtype=dtype,
            trust_remote_code=True,
        ).to(device).eval()

    @torch.no_grad()
    def predict(
        self,
        image: Image.Image,
        question: str,
        generation_mode: str = "hybrid",   # "fast" (MTP) | "slow" (NTP/AR) | "hybrid"
        max_new_tokens: int = 2048,
        temperature: float = 0.7,
        verbose: bool = True,
    ) -> dict:
        messages = [
            {"role": "user", "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": question},
            ]}
        ]

        text = self.processor.py_apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        images, videos = self.processor.process_vision_info(messages)
        inputs = self.processor(
            text=[text], images=images, videos=videos, return_tensors="pt"
        ).to(self.device)

        pixel_values = inputs["pixel_values"].to(self.dtype)
        input_ids = inputs["input_ids"]
        image_grid_hws = inputs.get("image_grid_hws", None)

        response = self.model.generate(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=inputs["attention_mask"],
            image_grid_hws=image_grid_hws,
            tokenizer=self.tokenizer,
            max_new_tokens=max_new_tokens,
            use_cache=True,
            generation_mode=generation_mode,
            temperature=temperature,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.1,
            verbose=verbose,
        )

        result = {"answer": response[0] if isinstance(response, tuple) else response}
        if isinstance(response, tuple) and len(response) >= 3:
            result["history"] = response[1]
            result["stats"] = response[2]
        return result

    # ---- Convenience methods for each task ----

    def detect(self, image: Image.Image, categories: list[str], **kwargs) -> dict:
        """Object detection / document layout analysis."""
        cats = "</c>".join(categories)
        prompt = f"Locate all the instances that matches the following description: {cats}."
        return self.predict(image, prompt, **kwargs)

    def ground_single(self, image: Image.Image, phrase: str, **kwargs) -> dict:
        """Phrase grounding — single instance."""
        prompt = f"Locate a single instance that matches the following description: {phrase}."
        return self.predict(image, prompt, **kwargs)

    def ground_multi(self, image: Image.Image, phrase: str, **kwargs) -> dict:
        """Phrase grounding — multiple instances."""
        prompt = f"Locate all the instances that match the following description: {phrase}."
        return self.predict(image, prompt, **kwargs)

    def ground_text(self, image: Image.Image, phrase: str, **kwargs) -> dict:
        """Text grounding."""
        prompt = f"Please locate the text referred as {phrase}."
        return self.predict(image, prompt, **kwargs)

    def detect_text(self, image: Image.Image, **kwargs) -> dict:
        """Scene text detection."""
        prompt = "Detect all the text in box format."
        return self.predict(image, prompt, **kwargs)

    def ground_gui(self, image: Image.Image, phrase: str, output_type: str = "box", **kwargs) -> dict:
        """GUI grounding (box or point)."""
        if output_type == "point":
            prompt = f"Point to: {phrase}."
        else:
            prompt = f"Locate the region that matches the following description: {phrase}."
        return self.predict(image, prompt, **kwargs)

    def point(self, image: Image.Image, phrase: str, **kwargs) -> dict:
        """Pointing."""
        prompt = f"Point to: {phrase}."
        return self.predict(image, prompt, **kwargs)

    # ---- Utility: parse model output ----

    @staticmethod
    def parse_boxes(answer: str, image_width: int, image_height: int) -> list[dict]:
        """Parse model output into pixel-coordinate bounding boxes.

        Coordinates in model output are normalized integers in [0, 1000].
        """
        boxes = []
        for m in re.finditer(r"<box><(\d+)><(\d+)><(\d+)><(\d+)></box>", answer):
            x1, y1, x2, y2 = [int(g) for g in m.groups()]
            boxes.append({
                "x1": x1 / 1000 * image_width,
                "y1": y1 / 1000 * image_height,
                "x2": x2 / 1000 * image_width,
                "y2": y2 / 1000 * image_height,
            })
        return boxes

    @staticmethod
    def parse_points(answer: str, image_width: int, image_height: int) -> list[dict]:
        """Parse model output into pixel-coordinate points."""
        points = []
        for m in re.finditer(r"<box><(\d+)><(\d+)></box>", answer):
            x, y = int(m.group(1)), int(m.group(2))
            points.append({
                "x": x / 1000 * image_width,
                "y": y / 1000 * image_height,
            })
        return points

worker = LocateAnythingWorker("nvidia/LocateAnything-3B")
img = Image.open("example.jpg").convert("RGB")

# Scene Text Detection
result = worker.detect_text(img)
print("Text Detection:", result["answer"])


# Parse structured output into pixel coordinates
w, h = img.size
boxes = LocateAnythingWorker.parse_boxes(result["answer"], w, h)
points = LocateAnythingWorker.parse_points(result["answer"], w, h)

[michaelmanly]

Thanks, that helps. Since this is nvidia/LocateAnything-3B and the failure is trying to allocate ~40GiB during generate(), I’d first try reducing the peak generation/memory pressure before assuming the model itself cannot run on L40.

A few things worth testing:

result = worker.detect_text(
    img,
    max_new_tokens=512,        # or 256 for a smoke test
    generation_mode="fast",    # compare against hybrid
    temperature=0.0,
    do_sample=False,
)

And before running:

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

Also worth checking the input image resolution, because the vision tokens can dominate memory here. A smaller resized image may confirm whether this is input-size driven.

If it still OOMs after lowering max_new_tokens and image size, then this is probably a real VRAM ceiling for that inference path on a single L40. In that case the clean test would be to run the same script once on an 80GB GPU with a hard cost/runtime cap, just to confirm whether it’s capacity rather than environment/config.

For example:

npm install -g badgr-cli
badgr login
badgr run python locate_anything_test.py --gpu A100 --max-cost 5 --max-runtime 30

That would make it a bounded test instead of leaving a larger GPU session running while debugging.
@dianjingshilezhi

[jsshizhan]

same issue

[exiawsh]

@dianjingshilezhi @jsshizhan @michaelmanly
Thank you to the community for your interest in our work! We now support pure Flash Attention for inference, which means you can run inference with lower VRAM usage and faster speeds on GPUs like the RTX 4090 and A100. Please refer to our latest README: nvidia/LocateAnything-3B · Hugging Face. We will also update our GitHub repository tomorrow with more detailed instructions!

[lixiaokun777]

Has the same issue been resolved

[jsshizhan]

@exiawsh The issue persists on the NVIDIA GeForce RTX 5090

[exiawsh]

@jsshizhan
To better understand and resolve this OutOfMemoryError, could you provide a bit more detail?

Have you tried running this using la_flash?

Is flash-attention successfully installed in your environment?

What is the resolution of the input image you are using?

Without these specific details, it is difficult to accurately pinpoint the cause of the issue and provide a fix. Let us know!