A Hidden 10% CPU Cost in TensorFlow

While debugging a TensorFlow performance issue with a colleague, Vadim, we uncovered a surprising source of overhead that had been hiding in plain sight. The problem wasn't in the model code or the GPU pipeline—it was in the memory allocator's interaction with the kernel, costing a significant chunk of CPU time.

Spotting the Anomaly

Vadim had flagged an unusual pattern in a CPU flame graph: a tall orange tower representing kernel code that stood out from the rest of the stacks. Closer inspection revealed the cause—10% of total CPU time was being spent in page faults.

At scale, an unexpected 10% CPU cost across thousands of server instances is a major financial and operational issue. This is exactly the kind of problem where flame graph analysis pays off: chasing down even single-percent inefficiencies can yield substantial savings.

Why Continuous Faulting?

The stack traces showed the faults originated from __memcpy_avx_unaligned(), which makes sense for a data segment mapping. However, the process had been running for hours, and by that point most mappings should have been faulted in, with page faults tapering off. Something was resetting those mappings.

One common cause is frequent mmap()/munmap() calls, but tracing with eBPF tools like mmapsnoop.py showed no such activity. The next suspect was madvise() dropping memory, and the flame graph confirmed it: madvise() accounted for 0.8% of CPU time, with the kernel's zap_page_range() path doing the heavy lifting.

The Premature Optimization Problem

Reading the kernel source in mm/advise.c showed that madvise() with the MADV_DONTNEED flag calls zap_page_range(). This flag tells the kernel the pages are no longer needed, dropping the virtual-to-physical mappings. The allocator was using MADV_DONTNEED on pages it actually needed, forcing the kernel to tear down mappings that would soon have to be rebuilt—triggering a fresh page fault each time.

This looked like a textbook case of premature optimization gone wrong: the allocator was discarding mappings to signal memory availability, but the cost of re-faulting those pages was far higher than any benefit gained.

Allocator Swap Fixes the Issue

Suspecting an allocator-level problem, Vadim checked the TensorFlow build configuration. The culprit was jemalloc, which was being used as the memory allocator. Rebuilding with glibc's allocator eliminated the problem entirely.

The corrected flame graph showed no more page fault tower, and initial testing measured a 3% performance improvement from the swap—not the 10% one might expect from the fault overhead, but still a meaningful win. The real cost was likely masked by overlapping execution, but the eliminination of avoidable kernel work was a clear improvement.