- Posted on
- • Artificial Intelligence
Artificial Intelligence Memory Analysis
- Author
-
-
- User
- linuxbash
- Posts by this author
- Posts by this author
-
Artificial Intelligence Memory Analysis on Linux: A Bash-first Playbook
If you’ve ever watched a training job crash with “Out of memory” at 97% progress, you know the pain. AI workloads burn through RAM and VRAM in unpredictable ways: data pipelines bloat page cache, libraries map enormous anonymous regions, and GPU allocators fragment VRAM. The good news: with a small set of Linux-native tools and a few Bash one-liners, you can turn “mysterious OOM” into actionable insight.
This article gives you a pragmatic, Bash-centric workflow for analyzing and controlling memory in AI workloads on Linux—covering what to measure, how to measure it, and how to prevent future blowups. You’ll get ready-to-run commands, install instructions for apt, dnf, and zypper, and battle-tested tips for both CPU RAM and GPU VRAM.
Why AI memory analysis matters
Reproducibility: Two runs, same code, different batches—one fits, one OOMs. Understanding RSS, PSS, and page cache makes behavior repeatable.
Performance: Memory pressure increases minor page faults, swap, and I/O stalls; you pay in epoch time.
Cost: Every GB of RAM/VRAM you can reclaim increases batch size and throughput or shrinks your instance size.
Stability: Catch leaks, fragmentation, and unbounded buffers before they crash your long runs.
Quick install: tools we’ll use
Run these with sudo as needed. Choose the block for your distro’s package manager.
APT (Debian/Ubuntu)
sudo apt update
sudo apt install -y smem htop sysstat nvtop valgrind numactl
DNF (Fedora/RHEL/CentOS Stream)
sudo dnf install -y smem htop sysstat nvtop valgrind numactl
Zypper (openSUSE)
sudo zypper refresh
sudo zypper install -y smem htop sysstat nvtop valgrind numactl
Optional (Python tools)
python3 -m pip install --upgrade pip
python3 -m pip install memray
1) Get the right numbers: RSS vs PSS vs VSS (and why your “memory use” is lying)
RSS (Resident Set Size): actual physical memory used by a process. Can double-count shared pages across processes.
PSS (Proportional Set Size): like RSS but shared pages are divided across processes. Best single metric for “who really uses RAM.”
VSS (Virtual Size): address space mapped. Often huge and misleading for real usage.
Commands you’ll actually use:
- Quick system view
free -h
grep -E 'MemTotal|MemFree|MemAvailable|Buffers|Cached' /proc/meminfo
- Per-process PSS (requires smem)
smem -c "pid user pss rss command" | sort -n -k3 | tail -20
- Fast per-PID summary (kernel smaps_rollup)
pid=<YOUR_PID>
sudo awk '/^(Pss|Rss|Swap):/ {sum[$1]+=$2} END {for (k in sum) printf "%s %d kB\n", k, sum[k]}' /proc/$pid/smaps_rollup
- Detailed mappings (useful for hunting huge mmaps)
pmap -x <YOUR_PID> | tail -n 1
- TUI for quick sanity checks
htop
# Press F2 -> Columns -> add PSS (if available), RES, SHR; save layout
Real-world example:
- Your Python training shows “only 8 GB RSS” but the system is out of memory. smem reveals PSS is actually 28 GB because many workers map the same 20 GB of dataset files; RSS was double-counting shared cache pages. PSS exposes the true load.
2) Log memory over time (and catch OOMs in the act)
Short runs look fine; long runs leak. Sample memory and the kernel’s OOM messages, then correlate with your training timeline.
- Lightweight process logger (RSS, PSS, and system stats)
#!/usr/bin/env bash
# follow-mem.sh PID INTERVAL
set -euo pipefail
pid="${1:?PID required}"
int="${2:-2}"
echo "ts,vmrss_kb,pss_kb,swap_kb,mem_avail_kb,page_cache_kb"
while [ -d "/proc/$pid" ]; do
ts=$(date +%s)
vmrss=$(awk '/VmRSS:/ {print $2}' /proc/$pid/status 2>/dev/null || echo 0)
read -r pss swap < <(sudo awk '/^(Pss|Swap):/ {gsub(":",""); a[$1]+=$2} END {print a["Pss"], a["Swap"]}' /proc/$pid/smaps_rollup 2>/dev/null || echo "0 0")
mem_avail=$(awk '/MemAvailable:/ {print $2}' /proc/meminfo)
page_cache=$(awk '/^Cached:/ {print $2}' /proc/meminfo)
echo "$ts,$vmrss,$pss,$swap,$mem_avail,$page_cache"
sleep "$int"
done
- Kernel OOM messages
journalctl -k -S today | grep -Ei "out of memory|Killed process" -n
- System-wide RAM trend
sar -r 1 60 # 1-second interval for 60 samples
What to look for:
PSS rising linearly per iteration → leak or unbounded buffer.
Drops in MemAvailable with rising Cached → data pipeline hammering page cache.
Swap growth → thrashing; expect massive slowdown before OOM.
3) Page cache, datasets, and /dev/shm: tame the hidden consumers
AI pipelines often stress the filesystem cache and shared memory:
- Measure file cache pressure
free -h
grep -E 'Cached|Buffers|Dirty|Writeback' /proc/meminfo
- Carefully drop caches during experiments (not in production)
sudo sync
echo 3 | sudo tee /proc/sys/vm/drop_caches
- Check and resize /dev/shm (important for DDP, dataloaders, multiprocessing queues)
df -h /dev/shm
sudo mount -o remount,size=32G /dev/shm
- Dataset workers and prefetching (PyTorch example)
- Fewer workers, smaller prefetch
- Use memory-mapped datasets judiciously; monitor PSS and page cache
- Example code change:
DataLoader(dataset, num_workers=2, prefetch_factor=2, persistent_workers=True, pin_memory=True)
Real-world example:
- OOMs only on large-image datasets? PSS shows moderate growth, but MemAvailable plummets while Cached grows. Reducing prefetch_factor and remounting /dev/shm avoids cache spikes and stabilizes the run.
4) Contain the blast radius with cgroups (systemd) and containers
Prevent one runaway job from taking down the host by enforcing memory ceilings.
- Run under a memory cap with systemd-run (root)
sudo systemd-run --scope -p MemoryMax=40G bash -lc 'python train.py --config config.yaml'
- Same idea with Docker
docker run --rm --memory=40g --memory-swap=40g \
-v $PWD:/workspace -w /workspace myimage:latest \
python train.py
- NUMA-aware allocations (can reduce cross-socket traffic and surprises)
numactl --cpunodebind=0 --membind=0 python train.py
Tips:
Set MemoryMax slightly below system RAM to leave headroom for the OS and I/O.
For multi-tenant boxes, put each experiment in its own scope/container with explicit limits.
5) Debug leaks and fragmentation: Python, C/C++, and GPU notes
- Python: trace high-level allocs (CPU RAM)
python -X tracemalloc -c "import tracemalloc; tracemalloc.start(); import your_script"
# Or inside code:
import tracemalloc, time
tracemalloc.start()
# ... training loop ...
print(tracemalloc.get_traced_memory())
- Deep inspection with Memray (pip install memray)
python -m memray run -o memray.bin train.py
python -m memray flamegraph memray.bin
- C/C++ inference servers: heap growth? Use Valgrind Massif
valgrind --tool=massif --massif-out-file=massif.out ./your_binary
ms_print massif.out | less
- GPU VRAM quick check across vendors: nvtop
nvtop
# Observe per-process VRAM and utilization in real time
- PyTorch VRAM fragmentation tip
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
# Inside Python:
import torch
print(torch.cuda.memory_summary())
What to look for:
Consistent increase per-epoch in tracemalloc or Memray → leak in preprocessing, caching, or model snapshotting.
Massif shows large retained heap from a specific code path → free or reuse buffers; consider arena/allocator tuning.
VRAM “free but OOM” → fragmentation; tweak allocator config or reduce peak temporary tensors.
Putting it together: a 10-minute triage
1) Install tools (smem, htop, sysstat, nvtop, valgrind, numactl). 2) Start your job under a cap:
sudo systemd-run --scope -p MemoryMax=40G bash -lc 'python train.py'
3) In another terminal:
htop
smem -c "pid user pss rss command" | sort -nk3 | tail -20
./follow-mem.sh <PID> 2 | tee mem.csv
4) If OOMs happen:
journalctl -k -S today | grep -Ei "out of memory|Killed process"
5) If cache ballooning is suspected:
grep -E 'Cached|Dirty|Writeback' /proc/meminfo
6) If long runs grow steadily, add tracemalloc/Memray or Massif and re-run for a shorter window.
Conclusion and next steps (CTA)
Memory issues in AI aren’t mysterious once you measure the right things. Start using PSS, log trends over time, control page cache and /dev/shm, and enforce sane cgroup limits. With these Bash-first techniques, you’ll convert late-epoch OOMs into predictable, stable runs—and often unlock bigger batches and faster throughput.
Your next step:
Install the tools above.
Wrap your next experiment with systemd-run and follow-mem.sh.
If you still see growth, instrument with Memray or Massif and fix the hotspot.
If you want a ready-made kit, turn the snippets here into a tiny “ai-mem-toolbox” repo for your team—one script to launch runs with caps, one to log memory, and a README of the common checks.