-
Notifications
You must be signed in to change notification settings - Fork 0
AI Performance
collective aggregation operations like all-reduce, all-to-all, and all-gather, which are extensively during model training and inference
-
Measure goodput: Look beyond raw FLOPS or utilization. Instead, measure the ratio of time the GPU spends performing useful work (e.g., forward/backprop computations) versus waiting on data or other overhead. Use NVIDIA Nsight Systems/Compute or PyTorch profiler to measure this ratio. Strive to improve this ratio, as goodput focuses on effective, useful GPU utilization.
-
Prefer skillful engineering optimizations instead of brute-force spending: More hardware isn’t a silver bullet. Clever software and system optimizations can bridge the gap when hardware is limited, enabling results that would otherwise require far more expensive infrastructure. We saw this with DeepSeek’s achievement. By optimizing communication and hardware usage, their engineers trained a frontier model on restricted H800 GPUs (with limited interconnect bandwidth) at a fraction of the cost. The DeepSeek model matched the performance of frontier models trained on far more powerful hardware. In other words, skillful engineering outperformed brute-force spending.
-
Look for order-of-magnitude impact with incremental optimizations: At scale, even a small-percentage efficiency gain can save millions of dollars. Said differently, small inefficiencies such as redundant computations and slow data pipelines can silently increase costs as the system scales.
-
Approach performance tuning with a profile-driven mindset: Use data and profiling tools to guide optimizations. Use profilers to identify the true bottlenecks—whether it’s compute utilization, memory bandwidth, memory latency, cache misses, or communication/network delays. Then apply targeted optimizations for that bottleneck.
-
Maintain a holistic view: Improving AI systems performance spans hardware, including the GPU, CPU, memory, and network—as well as software such as algorithms and libraries. A weakness in any layer can bottleneck the whole. The best performance engineers consider hardware-software codesign: sometimes algorithm changes can alleviate hardware limits, and sometimes new hardware features enable new algorithms.
-
Stay informed on the latest hardware, software, and algorithms: Modern AI hardware and software are evolving rapidly. New capabilities such as unified CPU-GPU memory, faster interconnects, and novel numerical-precision formats can change the optimal strategies. A good performance engineer keeps an eye on these and updates their mental models accordingly to eliminate bottlenecks quickly. Additionally, the MLPerf benchmark suites are a great resource to understand AI hardware performance for various models.
- cuPyNumeric: dropin replacement for numpy.
When a researcher submits a training job, the scheduler reserves nodes, the OS provides the GPU devices and memory allocations using the NVIDIA driver, and the container provides the correct software environment (including the optimized, hardware-aware CUDA libraries). The user code (e.g., PyTorch, TensorFlow, JAX) uses these CUDA libraries, which ultimately communicate with the driver and hardware.
A NUMA node is a logical grouping of CPUs, GPUs, network interface con‐ trollers (NICs), and memory that are physically close to one another.
By binding a process to a CPU on the same NUMA node as its GPU, we can avoid this extra overhead. For instance, you can use numactl --cpunodebind= --membind= to bind both CPU threads and memory allocations to the GPU’s local NUMA node.
The key idea is to keep CPU execution and memory access local to the GPU that it’s serving
By default, processes may be migrated across NUMA nodes. This will lead to additional latency caused by remote memory accesses. As such, it’s important to explicitly bind processes and memory to the same NUMA node as the local GPU. You can do this using numactl, taskset, or cgroups
To explicitly specify NUMA-affinity, you need to “pin” processes or threads to spe‐ cific CPUs that are connected to the same NUMA node as the GPU. This type of CPU affinity is called CPU pinning.
If you launch eight training processes, one per GPU, you should bind each training process to a CPU core—or set of CPU cores—connected to the same NUMA node as the GPUs
Linux provides tools to do this, including numactl --cpunodebind= --membind= , which launches a process pinned to the given NUMA node. You can also use taskset to pin processes to specific core IDs. Here is an example using numactl to bind the train.py script to a CPU running in the same NUMA
node 1 as GPU 4:
numactl --cpunodebind=1 --membind=1 \
python train.py --gpu 4This assumes we know the NUMA node ID and that we are binding the script to only one GPU. Binding the train.py to multiple GPUs to an unknown NUMA node is a bit more complicated. The following script dynamically queries the topology using nvidia-smi topo and binds the script to GPUs using the local NUMA node:
#!/bin/bash
for GPU in 0 1 2 3; do
# Query NUMA node for this GPU
NODE=$(nvidia-smi topo -m -i $GPU \
| awk '/NUMA Affinity/ {print $NF}')
# Launch the training process pinned to that NUMA node
numactl --cpunodebind=$NODE --membind=$NODE \
bash -c "CUDA_VISIBLE_DEVICES=$GPU python train.py --gpu $GPU"
doneMany deep learning frameworks also let you set thread affinities programmatically. For instance, PyTorch’s DataLoader exposes worker_init_fn so you can set CPU affinity for each worker process during initialization
import os
import re
import glob
import subprocess
import psutil
import ctypes
import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, Dataset
from functools import partial
# Optional: NVML is preferred for GPU↔NUMA mapping
try:
import pynvml as nvml # pip install nvidia-ml-py3
_HAS_NVML = True
except Exception:
_HAS_NVML = False
# --- libnuma for memory binding
_libnuma = ctypes.CDLL("libnuma.so")
if _libnuma.numa_available() < 0:
raise RuntimeError("NUMA not available on this system")
_libnuma.numa_run_on_node.argtypes = [ctypes.c_int]
_libnuma.numa_set_preferred.argtypes = [ctypes.c_int]
def parse_physical_cpu_list(phys_str: str):
"""Parse '0-3,8-11' -> [0,1,2,3,8,9,10,11]."""
cpus = []
if not phys_str:
return cpus
for part in phys_str.split(','):
part = part.strip()
if not part:
continue
if '-' in part:
start, end = map(int, part.split('-'))
cpus.extend(range(start, end + 1))
else:
cpus.append(int(part))
return cpus
def get_numa_cpus_for_node(node: int):
"""Read /sys/devices/system/node/node{node}/cpulist."""
path = f"/sys/devices/system/node/node{node}/cpulist"
with open(path, "r") as f:
return parse_physical_cpu_list(f.read().strip())
def get_numa_cpus_and_memory():
"""Return (current_cpu_mask, preferred_node) from numactl --show."""
out = subprocess.run(
["numactl", "--show"],
capture_output=True,
text=True
).stdout
phys = re.search(r"physcpubind:\s*([\d,\-\s]+)", out).group(1)
cpus = parse_physical_cpu_list(phys)
node = int(
re.search(r"preferred node:\s*(-?\d+)", out).group(1)
)
return cpus, node
def get_gpu_numa_node(device: int) -> int:
"""
Determine NUMA node for a GPU (prefer NVML; fall back to sysfs;
final fallback to current preferred node).
"""
# NVML path (preferred)
if _HAS_NVML:
try:
nvml.nvmlInit()
props = torch.cuda.get_device_properties(device)
pci = props.pci_bus_id
# Normalize to 8-hex-digit domain if needed for NVML
try:
domain, bus, devfn = pci.split(':')
if len(domain) < 8:
domain = domain.rjust(8, '0')
pci8 = f"{domain}:{bus}:{devfn}"
except ValueError:
pci8 = pci
try:
handle = nvml.nvmlDeviceGetHandleByPciBusId_v2(pci8)
except AttributeError:
handle = nvml.nvmlDeviceGetHandleByPciBusId(pci8)
# Direct NUMA ID if driver exposes it
try:
numa_id = nvml.nvmlDeviceGetNUMANodeId(handle)
if isinstance(numa_id, int) and numa_id >= 0:
return numa_id
except Exception:
pass
# Derive from NVML CPU affinity
cpu_count = psutil.cpu_count(logical=True)
elems = (cpu_count + 63) // 64
mask = nvml.nvmlDeviceGetCpuAffinity(handle, elems)
cpus = []
for i, m in enumerate(mask):
m = int(m)
for b in range(64):
if m & (1 << b):
cpu_id = i * 64 + b
if cpu_id < cpu_count:
cpus.append(cpu_id)
# Build CPU→NUMA map from sysfs and choose majority node
cpu2node = {}
for node_path in sorted(
glob.glob("/sys/devices/system/node/node*")
):
node_id = int(
os.path.basename(node_path).replace("node", "")
)
with open(os.path.join(node_path, "cpulist"), "r") as f:
for c in parse_physical_cpu_list(f.read().strip()):
cpu2node[c] = node_id
counts = {}
for c in cpus:
n = cpu2node.get(c)
if n is not None:
counts[n] = counts.get(n, 0) + 1
if counts:
return max(
counts.items(),
key=lambda kv: kv[1]
)[0]
except Exception:
pass
# sysfs fallback
try:
props = torch.cuda.get_device_properties(device)
pci = props.pci_bus_id
sysfs_path = f"/sys/bus/pci/devices/{pci}/numa_node"
with open(sysfs_path, "r") as f:
val = int(f.read().strip())
return val if val >= 0 else 0
except Exception:
pass
# Last resort: current preferred node
_, node = get_numa_cpus_and_memory()
return node if node >= 0 else 0
def set_numa_affinity(node: int):
"""Bind current process to CPUs and memory of the given NUMA node."""
cpus = get_numa_cpus_for_node(node)
# IMPORTANT: CPUs of target node
psutil.Process(os.getpid()).cpu_affinity(cpus)
_libnuma.numa_run_on_node(node)
_libnuma.numa_set_preferred(node)
print(
f"PID={os.getpid()} bound to NUMA node {node} "
f"(CPUs={cpus})"
)
return cpus
def _worker_init_fn(worker_id: int, node: int, cpus: list):
"""Reapply binding in each DataLoader worker (no CUDA calls here)."""
psutil.Process(os.getpid()).cpu_affinity(cpus)
_libnuma.numa_run_on_node(node)
_libnuma.numa_set_preferred(node)
print(
f"Worker {worker_id} "
f"(PID={os.getpid()}) "
f"bound to NUMA node {node}"
)
# ----- Example usage below -----
class MyDataset(Dataset):
def __len__(self):
return 1024
def __getitem__(self, idx):
return torch.randn(224 * 224 * 3, device="cpu")
def main():
# DDP setup
dist.init_process_group(
backend="nccl",
init_method="env://"
)
device = torch.cuda.current_device()
# Determine GPU's NUMA node and bind this process
gpu_node = get_gpu_numa_node(device)
cpus = set_numa_affinity(gpu_node)
# Build dataloader with closure-based worker_init_fn
dataset = MyDataset()
init_fn = partial(
_worker_init_fn,
node=gpu_node,
cpus=cpus
)
dataloader = DataLoader(
dataset,
batch_size=32,
num_workers=4,
pin_memory=True,
persistent_workers=True,
worker_init_fn=init_fn,
prefetch_factor=2,
)
# Model and DDP
model = torch.nn.Linear(
224 * 224 * 3,
10,
bias=True
).to("cuda")
ddp_model = DDP(
model,
device_ids=[device],
static_graph=True
)
for batch in dataloader:
batch = batch.to(
"cuda",
non_blocking=True
)
out = ddp_model(batch)
# ... loss, backward, optimizer ...
if __name__ == "__main__":
main()
This script binds the main training process and each DataLoader worker process to the GPU’s local NUMA node to prevent cross-NUMA memory access. In the DataLoader, we pass a closure-based worker_init_fn that reapplies the precomputed NUMA binding inside each worker. And we do this without touching any CUDA APIs in the worker.
Because some launchers, container runtimes, or kernels do not reliably propagate NUMA policy to children, we explicitly reapply and verify the binding in every forked worker. It’s not safe to rely on inheritance alone.
Remember to set pin_memory=True and use non_blocking=True on H2D copies so that page-locked host buffers stay on the correct NUMA node. Prefer persis tent_workers=True to avoid re-forking workers and losing their affinity between epochs. And do not call torch.cuda.* in worker_init_fn. Instead, pass the GPU index using a closure or environment variable. The result is that data preparation and batch loading happen entirely in local mem‐ ory. This way, your GPUs stay busy and never need to pause for a remote‐NUMA hop. With this code, you get robust, topology‐aware affinity on any Linux server with libnuma and numactl installe
By default, numactl applies its CPU and memory policy to a process and is docu‐ mented to inherit that policy to all forked children. In practice, however, threads spawned by Python frameworks or exec’d subprocesses don’t always pick up the same settings on every kernel or Linux distribution. When using framework-managed worker processes, you should explicitly reassert the CPU and memory policy inside of each worker.
With a superchip architecture like Grace Blackwell (and Vera Rubin), the CPU and GPU are coherent using NVLink-C2C. How‐ ever, Linux still models CPU DRAM and GPU HBM as separate pools. Binding CPU threads to the local CPU NUMA node still remains beneficial for locality
In practice, pinning can eliminate unpredictable CPU scheduling behavior. It ensures that a critical thread such as a data-loading thread for your GPU doesn’t suddenly get migrated by the OS to a core on a different NUMA node in the middle of training or inferencing. In practice, it’s possible to see 5%–10% training throughput improve‐ ments just by eliminating cross-NUMA traffic and CPU core migrations. This also tends to reduce performance jitter and variance.
Many high-performance AI systems evaluate CPU simultaneous multithreading (SMT), or hyperthreading as it’s often called—and sometimes disable it for more pre‐ dictable per-core performance, but the benefit is workload-dependent. These systems may also reserve a handful of cores exclusively for OS background tasks by setting the isolcpus kernel parameter to isolate them from the general scheduler. You can also use Kubernetes CPU isolation for system daemons. This ensures that the remaining cores are dedicated entirely to training and inference threads and doing useful work.
By default, a process will allocate memory from the NUMA node of the CPU it’s cur‐ rently running on. So if you pin a process to NUMA node 0, its memory will natu‐ rally come from NUMA node 0’s local RAM, which is ideal. However, if the OS scheduler migrates threads—or if some memory was allocated before you did the pinning—you could end up with the nonideal scenario in which a process running in NUMA node 0 is using memory from NUMA node 1. To avoid this, the numactl --membind option forces memory allocation from a spe‐ cific NUMA node The general rule is to keep memory close to the CPU, which is close to the GPU.
As such, any worker subprocesses forked by your training script will automatically use the same NUMA memory binding. How‐ ever, they must be created using a fork-based model. If you switch to a spawn start method, or otherwise exec a new program, those child processes do not inherit the parent’s memory policy.
In addition, pinned memory, also called page-locked memory, is essential for efficient and direct GPU access. When memory is pinned, the OS won’t swap or move it. This leads to faster direct memory access (DMA) transfers.
You can test the data-transfer bandwidth between CPU memory and GPU memory using bandwidthTest --memory= from the installed CUDA utilities.
Deep learning frameworks provide options to use pinned memory for data loaders. For example, PyTorch’s DataLoader has a flag pin_memory=True
n short, you should make sure that your data loader uses pinned memory (e.g., pin_memory=True in PyTorch DataLoader) and that GPUDirect RDMA and GDS are enabled for supported hardware. This will reduce data transfer latency.
If you plan to use large, pinned buffers, ensure the ulimit value is high—or set it to unlimited. Otherwise the allocation might fail. Typically, one sets it to unlimited for large AI workloads and highperformance computing (HPC) applications.
Linux memory management typically uses 4 KB pages, but managing millions of tiny pages is inefficient when you have processes using tens or hundreds of gigabytes of memory, as in the case of deep learning data‐ sets, prefetched batches, model parameters, etc. Hugepages—2 MB or even 1 GB pages—can reduce the overhead of virtual memory management by making memory chunks bigger. The main benefits are fewer page faults and less pressure on the translation lookaside Buffer (TLB).
In scenarios with very large memory pools (e.g., preallocated pinned buffers for I/O), you may also consider explicit hugepages using vm.nr_hugepages or hugetlbfs for more deterministic performance.
Remember that, when using large, pinned memory regions, you should raise the ulimit -l setting (max locked memory) to a high value or unlimited. If this limit is too low, your attempt to pin memory can fail, leading to fallback on swappable memory—or out-of-memory (OOM) errors.
It’s important to note that THP’s background compaction can introduce unpredicta‐ ble pauses that are disastrous for latency-sensitive LLM inference workloads. Linux is configured by default to use THP to automatically allocate 2 MB pages whenever pos‐ sible. This is often sufficient, but it’s worth testing for your workload.
You can disable THP, but you will need to manually allocate and control hugepages. This will incur extra complexity, but it might be needed for low-latency workloads like inference. With THP disabled, your system will avoid stalls caused by kerneldriven defragmentations.
The modern consensus is to enable THP for most GPU-based training workloads in which throughput is important and to disable THP completely (transparent_hugepage=never)—or use madvise—for workloads like inference in which latency is important. This is also true for distributed training workloads in which many ranks (GPUs) allocate memory simultaneously.
On a busy system, you want to make sure that important threads such as datapipeline threads aren’t interrupted frequently. Linux by default uses the Completely Fair Scheduler (CFS) that works well for most cases.
But if you have a very latency-sensitive thread that feeds the GPU with data, for example, you could consider using real-time first in, first out (FIFO) or round-robin (RR) priority scheduling for that thread. This would make sure that the high-priority thread runs without being preempted by normal-priority threads. In practice, however, if you’ve pinned your threads to dedicated cores, you often don’t need to mess with real-time thread priorities.
Another option is to isolate cores or create separate CPU partitions to further reduce interruptions on these dedicated compute resources. To do this, you can use cset, kernel parameters like isolcpus and nohz_full, or cgroup cpuset isolation. With isolation, the OS scheduler leaves those CPU cores for you to use as you wish.
cgroup CPU and memory affinity is strongly recommended in production environments. Using these, each AI workload is isolated on its own physical cores and memory regions. This will prevent cross-workload contention and NUMA penalties. Tools like cpuset cgroups or container runtimes (docker --cpuset-cpus) should be used to enforce this.
You can assign each device’s hardware interrupts to cores on the same NUMA node. This will prevent cross-node interrupt handling that would otherwise incur extra latency and evict useful cache lines on a remote node.
In practice, performance-sensitive systems often disable the default irqbalance daemon or run it with bespoke rules. The other option is to manually set each interrupt’s affinity mask using /proc/irq/*/smp_affinity. By pinning every GPU and NIC interrupt to the nearest cores, you guarantee that those device interrupts are always serviced on the optimal NUMA node.
In short, the combination of dedicated cores, appropriate scheduling priorities, and NUMA-aware hardware interrupt bindings can help minimize jitter for data loading threads that are feeding the GPUs.
you should always try to avoid memory swapping to avoid seeing a catastrophic, multiple-orders-of-magnitude slowdown.
We recommend setting vm.swappiness=0, which tells Linux to avoid swapping except under extreme memory pressure. It effectively isolates your training job’s memory with cgroup limits to prevent any swapping.
You should use cgroups v2 through Docker or Kubernetes to pin memory and CPUs to the AI process. This will enforce NUMA affinity and no-swap policies in containerized environments.
You can also use sudo swapoff -a to temporarily disable all swap devices and files until the next reboot. Just make sure you have enough RAM for your workload—or put limits to prevent overcommit. Otherwise, the OOM killer may reap the process. Monitor swap usage using vmstat or free -m to make sure swap stays at zero.
A best practice for large training jobs is to write frequent checkpoints to disk in case you need to restart a failed job from a known good checkpoint. During checkpointing, however, huge bursts of data might fill up the OS page cache and cause stalls.
For storage, you can adjust vm.dirty_ratio and vm.dirty_background_ratio to tune the page-cache size for buffering writes. For example, with multi-GB check‐ points, using a higher dirty ratio lets the OS batch more data in RAM before flushing to disk. This will smooth out large checkpoint writes and reduce stalls in your train‐ ing loop.
Another option is to perform checkpointing in a separate thread. A more recent option in PyTorch is to write distributed checkpoint partitions from nodes across the cluster. In this case, the checkpoint partitions will be combined when the checkpoint is loaded after a failed-job restart.
In latency-sensitive training workflows, it’s best to bypass the page cache entirely. For example, open checkpoint files with O_DIRECT or use Linux’s io_uring for asynchro‐ nous I/O to avoid page-cache stalls. After writing each checkpoint, call posix_fadvise (fd, 0, 0, POSIX_FADV_DONTNEED) to immediately drop those pages from cache and prevent memory pressure on subsequent iterations.
By default, many compute nodes will run CPUs in a power-saving mode, which either downclocks a CPU or puts it to sleep when it’s idle this can ause extra latency when the system wakes the CPUs up again when new work arrives.
Configure the governor to performance. This can be done using cpupower frequency-set -g performance or in the Basic Input/Output System (BIOS).
Likewise, disabling deep C-states can keep cores from going into a low-power sleep state. C0 is active; everything above C0 represents a deeper state of sleep
In practice, many server BIOS/UEFI (Unified Extensible Firmware Interface) offer a high-performance profile that automatically sets the CPU governor to “Performance” and disables deep C-states
Bubbles are periods of time when the GPU is wait‐ ing for the CPU to resume data processing. By keeping the CPU ready, we reduce such hiccups. Many BIOSes for servers have a setting to disable C-states—or at least limit them
You should always turn off anything in your system that might introduce unpredictable latency, such as excess context switching, CPU frequency scaling, and memory-to-disk swapping. The result should be that your CPUs deliver data to the GPUs as fast as the GPUs can consume it, without the OS scheduling things on the wrong core or taking CPU cycles away at the wrong tim
On a well-tuned GPU server, CPU usage may not be very high since GPUs handle most of the computation. However, CPU usage should remain steady and in lockstep with GPU activity. The CPUs must stay busy preparing each incoming batch while the current batch is being processed by the GPU.
By tuning your host’s memory allocator (jemalloc or tcmalloc), you can eliminate unpre‐ dictable pauses in data preparation.
After tuning, you should see each GPU’s utilization hover near 100% and drop only at required synchronization barriers. The GPUs should never stall for data due to CPUside delays. With jemalloc, you can shard allocations into per-CPU arenas (narenas), enable background_thread for off-path purging, and lengthen dirty_decay_ms/muzzy_decay_ms so that freed pages aren’t immediately returned to the OS. This will minimize lock contention and fragmentation.
You can tune jemalloc with the MALLOC_CONF environment variable as follows:
export MALLOC_CONF="narenas:8,dirty_decay_ms:10000,muzzy_decay_ms:10000
,background_thread:true"
Similarly, tcmalloc benefits from tuning the TCMALLOC_MAX_TOTAL_THREAD_ CACHE_BYTES and TCMALLOC_RELEASE_RATE environment variables. These will provide larger per-thread caches so that small allocations avoid global locks and syscalls keeping CPU threads ready to feed the GPU with low, predictable latency. You can do this as follows:
export TCMALLOC_MAX_TOTAL_THREAD_CACHE_BYTES=$((512*1024*1024))
export TCMALLOC_RELEASE_RATE=16
In short, optimizing the allocator can reduce allocator overhead and fragmentation. This will keep CPU threads consistently fast and avoid unexpected stalls feeding the GPU. Experiment with these environment variables and tune them for your specific workload and environment
GPU persistence mode, the partitions of MPS, MIG, and a few other considerations like clock settings, ECC memory, and out-of-memory behavior.