Use the official PyTorch CUDA image (no installs)
1) Create a working directory
mkdir -p gpu-sanity && cd gpu-sanity
2) Create the Python check script
cat > gpu_sanity.py <<'PY'
import os, sys, time
import torch
def fail(msg):
print(f"FAIL: {msg}", file=sys.stderr)
sys.exit(1)
# -----------------------------------------------------------------------------#
# This program performs a minimal "sanity check" of CUDA by:
# (1) verifying that a CUDA device is visible,
# (2) allocating GPU tensors, and
# (3) executing a matrix multiply, then reporting time and memory.
#
# The intent is not benchmarking. It is an operational check: "can I launch
# kernels and move real data through the GPU?" In that sense it resembles the
# classic kernel probes: simple, direct, and designed to fail loudly.
# -----------------------------------------------------------------------------
def _use_color() -> bool:
# We follow the customary terminal convention: do not emit escape sequences
# unless the output is a terminal, and allow an explicit opt-out via NO_COLOR.
if os.environ.get("NO_COLOR"):
return False
try:
return sys.stdout.isatty()
except Exception:
return False
_COLOR = _use_color()
def style(text: str, *codes: str) -> str:
# ANSI SGR "Select Graphic Rendition". This is deliberately small-bore:
# a handful of numeric codes and a reset at the end.
if not _COLOR or not codes:
return text
return f"\x1b[{';'.join(codes)}m{text}\x1b[0m"
def banner(text: str) -> str:
# A mild "hipster" palette: teal + violet + pink accents.
left = style("===", "38;5;45", "1") # teal
mid = style(text, "38;5;141", "1") # violet
right = style("===", "38;5;205", "1") # pink
return f"{left} {mid} {right}"
def label(key: str) -> str:
return style(f"{key}:", "38;5;110", "1") # denim blue
def ok(text: str) -> str:
return style(text, "38;5;114", "1") # mint
def warn(text: str) -> str:
return style(text, "38;5;215", "1") # peach
def bad(text: str) -> str:
return style(text, "38;5;203", "1") # soft red
def bytes_fmt(n: int) -> str:
# Present byte counts in binary units. One does not improve correctness
# by printing large integers without context; the operator needs scale.
units = ["B", "KiB", "MiB", "GiB", "TiB"]
v = float(n)
for u in units:
if v < 1024.0 or u == units[-1]:
return f"{v:.2f} {u}"
v /= 1024.0
return f"{n} B"
def is_oom_exc(e: BaseException) -> bool:
# CUDA OOM can surface as torch.OutOfMemoryError or as a wrapper error
# type depending on where it is raised. We treat all "out of memory"
# conditions uniformly to support retry with a smaller problem size.
if isinstance(e, torch.OutOfMemoryError):
return True
if type(e).__name__ in {"AcceleratorError", "CudaError"}:
return "out of memory" in str(e).lower()
return False
def cuda_mem_info_safe(device: torch.device):
# mem_get_info is convenient but not always robust under driver/container
# mismatches or severe memory pressure. When it fails, we still report the
# device total and proceed with a conservative allocation strategy.
try:
free_b, total_b = torch.cuda.mem_get_info(device)
return free_b, total_b, None
except Exception as e:
props = torch.cuda.get_device_properties(device)
return None, int(props.total_memory), e
def pick_matmul_n(device: torch.device) -> int:
# The original fixed size (e.g., n=4096) is a fine idea on an idle device,
# but it fails the moment the GPU is already busy. Here we compute a size
# from the free memory (when available), and accept an explicit override.
env_n = os.environ.get("GPU_SANITY_N")
if env_n:
try:
return max(128, int(env_n))
except ValueError:
fail(f"GPU_SANITY_N must be an int, got {env_n!r}")
free_b, total_b, _ = cuda_mem_info_safe(device)
# Approx bytes:
# - a, b, c each n*n*2 bytes (fp16) => ~6*n^2 bytes
# - matmul workspace + allocator overhead => multiply by a safety factor
safety = float(os.environ.get("GPU_SANITY_SAFETY", "2.5"))
denom = 6.0 * safety
budget_b = int((free_b if free_b is not None else total_b * 0.25))
n = int((budget_b / denom) ** 0.5)
# Keep it reasonable for a "sanity check": fast, but non-trivial.
n = max(512, min(n, 8192))
return n
# Section heading: the program speaks in small, well-labeled facts.
print(banner("GPU Sanity Check (PyTorch)"))
print(label("torch.__version__"), torch.__version__)
print(label("cuda available"), ok("True") if torch.cuda.is_available() else bad("False"))
if not torch.cuda.is_available():
fail("torch.cuda.is_available() is False (no CUDA visible)")
print(label("device_count"), torch.cuda.device_count())
if torch.cuda.device_count() < 1:
fail("No CUDA devices detected")
print(label("device_name[0]"), torch.cuda.get_device_name(0))
print(label("CUDA_VISIBLE_DEVICES"), os.environ.get("CUDA_VISIBLE_DEVICES"))
# Deterministic-ish (as much as practical for a sanity check)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
# Matrix Multiplication on the GPU
device = torch.device("cuda:0")
torch.cuda.set_device(device)
torch.cuda.empty_cache()
props = torch.cuda.get_device_properties(device)
print(label("device_total_memory"), bytes_fmt(int(props.total_memory)))
print(label("device_capability"), f"{props.major}.{props.minor}")
try:
print(label("gpu_processes"))
print(style(torch.cuda.list_gpu_processes().rstrip(), "38;5;245"))
except Exception as e:
print(warn("WARN:"), f"torch.cuda.list_gpu_processes failed: {e}")
free_b, total_b, mem_err = cuda_mem_info_safe(device)
if free_b is None:
print(warn("WARN:"), f"torch.cuda.mem_get_info failed: {mem_err}")
print(label("cuda_mem"), f"total={bytes_fmt(total_b)} (free unknown)")
else:
print(label("cuda_mem"), f"free={bytes_fmt(free_b)} total={bytes_fmt(total_b)}")
# A tiny allocation is a canary. If this fails, there is little point
# continuing: either the GPU is genuinely exhausted or CUDA is in distress.
n_tiny = 1
try:
_ = torch.empty((n_tiny,), device=device, dtype=torch.float16)
except Exception as e:
if is_oom_exc(e):
fail(
"CUDA allocations are failing even for a tiny tensor; this usually means the GPU is fully occupied "
"by other processes or the container/driver stack doesn't fully support this GPU yet. "
"Check `nvidia-smi` for running processes, then retry."
)
raise
n = pick_matmul_n(device)
print(label("matmul_n"), f"{n} (override with GPU_SANITY_N)")
last_err = None
for attempt in range(6):
try:
# Allocate inputs. These are the principal consumers of memory here.
# We allocate them explicitly to make failure modes obvious.
a = torch.randn((n, n), device=device, dtype=torch.float16)
b = torch.randn((n, n), device=device, dtype=torch.float16)
break
except Exception as e:
if not is_oom_exc(e):
raise
last_err = e
torch.cuda.empty_cache()
if n <= 256:
fail(
"GPU is out of memory even for a small allocation (try closing other GPU workloads, "
"or set GPU_SANITY_N=128 and retry)."
)
n = max(256, n // 2)
print(warn("WARN:"), f"OOM allocating inputs; retrying with matmul_n={n}")
else:
raise last_err
torch.cuda.synchronize()
t0 = time.time()
# The actual work: a single GEMM. It is enough to demonstrate kernel launch,
# device execution, and the ability to bring a scalar result back to the host.
c = a @ b
torch.cuda.synchronize()
t1 = time.time()
# Touch result to ensure it's real
checksum = float(c[0, 0].item())
alloc = torch.cuda.memory_allocated()
max_alloc = torch.cuda.max_memory_allocated()
print(ok("MATMUL: succeeded"))
print(label("elapsed_sec"), f"{t1 - t0:.3f}")
print(label("checksum(c[0,0])"), f"{checksum}")
print(label("memory_allocated"), f"{alloc} bytes ({bytes_fmt(alloc)})")
print(label("max_memory_allocated"), f"{max_alloc} bytes ({bytes_fmt(max_alloc)})")
print(style("DELIVERABLE:", "38;5;141", "1"), "GPU matmul succeeded + memory stats visible.")
PY
3) Run it (no Dockerfile needed)
Pick a PyTorch image tag that matches your CUDA runtime. This example uses a common CUDA 12.1 runtime tag. If it doesn’t exist on your machine, try another *-cuda* tag.
docker run --rm -it --gpus all \
-v "$PWD:/work" -w /work \
pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime \
python gpu_sanity.py
Quick “must be true” preflight (host)
nvidia-smi
docker run --rm --gpus all nvidia/cuda:12.1.1-base-ubuntu22.04 nvidia-smi
Run it
#!/bin/bash
set -auexo pipefail
# DGX Spark GB10–friendly base with CUDA + PyTorch toolchain
## Source for base container: https://build.nvidia.com/spark/nemo-fine-tune/instructions
CTR="nvcr.io/nvidia/pytorch:25.08-py3"
# Other Option: pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime
docker run \
--rm -it \
--gpus all \
--ipc=host \
--ulimit memlock=-1 \
--ulimit stack=67108864 \
-v "$PWD:/work" -w /work \
$CTR \
python gpu_sanity.py
Common failure modes (fix now, don’t proceed)
-
torch.cuda.is_available() == Falseinside container:-
Host
nvidia-smifails → driver issue. -
Host works but container fails → install/configure nvidia-container-toolkit and restart Docker.
-
-
“no CUDA devices”:
- Missing
--gpus allor Docker is too old.
- Missing
-
Runs but is painfully slow:
- You’re on CPU (CUDA not visible) or using a CPU-only PyTorch image.
Once this is all working…
let’s do it by hand 
- drop into a shell in the PyTorch container
#!/bin/bash
set -auexo pipefail
# DGX Spark GB10–friendly base with CUDA + PyTorch toolchain
## Source for base container: https://build.nvidia.com/spark/nemo-fine-tune/instructions
CTR="nvcr.io/nvidia/pytorch:25.08-py3"
docker run \
--rm -it \
--gpus all \
--ipc=host \
--ulimit memlock=-1 \
--ulimit stack=67108864 \
-v "$PWD:/work" -w /work \
$CTR \
bash
- run
ipython -
In [1]: import torch; torch.__version__ Out[1]: '2.8.0a0+34c6371d24.nv25.08 -
In [2]: torch.cuda.is_available() Out[2]: True -
In [3]: torch.cuda.device_count() Out[3]: 1 -
In [4]: torch.cuda.get_device_name(0) Out[4]: 'NVIDIA GB10' -
In [9]: torch.cuda.get_device_properties(torch.device("cuda:0")) Out[9]: _CudaDeviceProperties(name='NVIDIA GB10', major=12, minor=1, total_memory=...) -
In [10]: torch.cuda.list_gpu_processes().rstrip() Out[10]: 'GPU:0\nno processes are running' -
In [12]: torch.cuda.mem_get_info(torch.device("cuda:0")) Out[12]: (..., ...) -
In [46]: n = 51200 -
In [47]: a = torch.randn((n, n), device=device, dtype=torch.float16) -
In [48]: b = torch.randn((n, n), device=device, dtype=torch.float16) -
In [49]: t1 = time.time(); torch.cuda.synchronize(); a @ b; torch.cuda.synchronize(); time.time() - t1 Out[49]: 7.9207212924957275 -
In [15]: torch.cuda.memory_allocated() Out[15]: 0 In [16]: torch.cuda.max_memory_allocated() Out[16]: 0
In PyTorch c = a @ b is matrix multiplication (the @ operator calls torch.matmul).
If a and b are 2D matrices with shapes (n, m) and (m, k), the result c has shape (n, k),
where: c[i, j] = sum_{t=0..m-1} a[i, t] * b[t, j]
In the script, a and b are both (n, n), so it computes an (n, n) product on the GPU (a
common “GEMM” / GEneral Matrix Multiply workload).
In Python, @ is the matrix-multiplication operator (added in PEP 465). The interpreter
resolves a @ b by calling special (“dunder”) methods:
- first tries
a.matmul(b) - if that returns
NotImplemented, triesb.rmatmul(a) - for in-place
a @= b, usesa.imatmul(b)
For PyTorch tensors, these are implemented on torch.Tensor, and ultimately dispatch to
PyTorch’s matmul implementation (equivalent to calling torch.matmul(a, b) for 2D tensors).
torch.cuda.memory_allocated() returns how many bytes of GPU memory are currently allocated by PyTorch for tensors on the current CUDA device.
- this measures PyTorch’s active allocations (tensors, buffers) - not necessarily total GPU
memory used by the process - PyTorch also keeps a caching allocator, so freed tensors may leave memory “reserved” for reuse; that won’t show up here. For that we can use
torch.cuda.memory_reserved()(andtorch.cuda.max_memory_reserved()). - in the script, we print it after the matmul to show roughly how much memory the test
actually consumed