PyTorch Memory Profiling in Production: A Guide to Efficiency

The Production Memory Paradox

The primary challenge with memory management in production is the discrepancy between peak usage and reserved memory. PyTorch uses a caching allocator to speed up GPU memory allocations. When a tensor is freed, the memory is not immediately returned to the system; instead, it is kept in a cache for future use. This leads to a common point of confusion: nvidia-smi might show 95% VRAM usage, while your actual allocated tensors only occupy 60%.

Memory Fragmentation in Long-Running Jobs

In a production environment, this gap is where memory fragmentation lives. Fragmentation occurs when the allocator has enough total free memory but cannot find a contiguous block large enough for a new request. This is particularly prevalent in workloads with dynamic batching or variable sequence lengths, such as LLM inference. In unoptimized clusters, fragmentation can account for a meaningful share of wasted VRAM.

• Internal Fragmentation: Memory wasted within a block because the requested size was slightly smaller than the block provided.
• External Fragmentation: Small gaps between allocated blocks that cannot be merged into a single large block.
• Silent OOMs: Errors that occur not because you lack memory, but because the allocator cannot defragment the cache fast enough.

To manage this, you must move beyond torch.cuda.memory_allocated(). While useful for a snapshot, it does not tell you the state of the cache. You need to monitor torch.cuda.memory_reserved() and compare it against the actual allocation to calculate your fragmentation ratio. High-performance teams use this ratio as a primary metric for triggering automated restarts or cache clears.

Lightweight Observability with Memory Snapshots

For production environments, torch.profiler is often too heavy. Enabling profile_memory=True and with_stack=True can significantly increase latency, making it unsuitable for continuous use. The alternative is memory snapshots. Introduced in recent PyTorch versions and refined in the 2025 releases, torch.cuda.memory._snapshot() provides a detailed view of every allocation and its associated Python stack trace with minimal overhead.

The beauty of snapshots lies in their 'flight recorder' capability. You can record a history of allocations in a circular buffer. When an Out-of-Memory (OOM) event occurs, you can dump this buffer to a file. This allows for a post-mortem analysis of exactly which operation caused the spike. PyTorch's own documentation puts Python trace collection at about 2 microseconds per trace and C++ frame collection at roughly 50 nanoseconds per frame, and says outright that you may consider enabling this on production jobs if you anticipate ever having to debug memory issues.

Implementation involves three steps:

1. Enable History: Call torch.cuda.memory._record_memory_history(max_entries=100000) at the start of your worker process. A bare True dispatches to the legacy signature, which keeps one trace entry and no allocation history.
2. Set Limits: Use max_entries to cap the ring buffer, which PyTorch otherwise defaults to sys.maxsize; each entry costs several KB, so a long-running worker has to set it.
3. Capture on Trigger: Wrap your main loop in a try-except block to catch torch.cuda.OutOfMemoryError and dump the snapshot.

Once captured, these snapshots can be uploaded to the PyTorch Memory Visualizer. This tool provides a timeline view of your VRAM, allowing you to see exactly how the caching allocator is splitting segments and where 'zombie' tensors are lingering in the cache.

Taming the Caching Allocator

If your profiling reveals high fragmentation, the solution usually lies in the PYTORCH_CUDA_ALLOC_CONF environment variable. This is the most powerful, yet underutilized, tool for production stability. By default, the allocator is optimized for speed, but you can tune it for memory density.

One critical setting is max_split_size_mb. This prevents the allocator from splitting large unused blocks into many small ones, which is a leading cause of fragmentation. For example, setting max_split_size_mb:512 ensures that large blocks remain intact, making them available for future large tensor requests. Properly tuned, this parameter can markedly reduce OOM errors in multi-tenant GPU environments.

Another advanced feature is expandable segments. When enabled via expandable_segments:True, PyTorch uses a different low-level allocation strategy that allows segments to grow without requiring contiguous physical memory. This leaves far fewer unusable slivers at the end of segments, which is where external fragmentation comes from. PyTorch still marks the option experimental, and it targets jobs whose allocation sizes change from one iteration to the next, such as batched inference with a varying batch size.

PyTorch's own guidance is to treat max_split_size_mb as a last resort, for a job that is aborting out of memory and showing a large amount of inactive split blocks in its memory_summary() output. Where the timeline instead shows persistent 'gaps' left by allocation sizes that shift between iterations, expandable_segments is the setting to reach for.

Automated Triggering and Post-Mortems

A robust production strategy does not wait for a crash. It uses threshold-based profiling. By monitoring torch.cuda.memory_reserved() via a background thread or a sidecar process, you can trigger a memory snapshot when usage exceeds a safe threshold, such as 90% of total capacity. This 'pre-OOM' snapshot is often more valuable than the one taken at the moment of failure, as it shows the state of the system leading up to the crisis.

Threshold-Based Profiling with Flight Recorder

In PyTorch 2.5, the introduction of the Flight Recorder for distributed jobs has further simplified this. While primarily designed for debugging stuck processes, it can be adapted to monitor memory health across a cluster. If one node in a Distributed Data Parallel (DDP) setup starts showing abnormal memory growth, the Flight Recorder can capture the state of all nodes simultaneously, helping you identify if the leak is due to a specific data shard or a desynchronized gradient update.

Continuous Profiling Integration

Common mistakes to avoid in production profiling:

• Continuous Profiling: Never leave torch.profiler running indefinitely. It will eventually exhaust host memory with trace data.
• Ignoring CPU Memory: GPU OOMs are often caused by the CPU being unable to feed the GPU fast enough, leading to a backlog of tensors in the input queue.
• Manual Cache Clearing: Avoid calling torch.cuda.empty_cache() in a tight loop. It forces a global synchronization and can destroy your throughput. Use it only between logical stages of a pipeline.

By integrating these triggers into your orchestration layer, you create a self-healing system. Lyceum's scheduling product adds memory and runtime prediction within a node, which helps catch jobs that will not fit before they fail.

Infrastructure-Level Optimization

While code-level profiling is essential, the underlying infrastructure plays a massive role in memory efficiency. In a sovereign European cloud environment, data sovereignty and performance must go hand-in-hand. Lyceum runs GPU compute in European data centres in Spain, Paris and the Nordics, billed per second with no base fee, and a 24/7 direct line to the tech team helps align your PyTorch configuration with the physical GPU architecture.

For instance, using torch.compile in PyTorch 2.x can significantly reduce memory usage through kernel fusion. By merging multiple operations into a single CUDA kernel, the system avoids materializing intermediate tensors in VRAM. On Lyceum's H100 and A100 instances you can validate torch.compile against your specific configuration, making sure the memory savings do not come at the cost of stability.

Ultimately, memory profiling is about predictability. In regulated industries like finance or healthcare, a production failure isn't just a metric; it's a compliance risk. By combining PyTorch's native snapshotting tools with sovereign, high-performance European infrastructure, you gain the visibility needed to scale AI workloads with confidence. You move from reactive firefighting to proactive resource management, ensuring that every byte of VRAM is contributing to model performance.

Sources

[1] PyTorch: Understanding CUDA Memory Usage, with the latency of _record_memory_history (read 4 August 2026); [2] PyTorch: Understanding GPU Memory 1, Visualizing All Allocations over Time; [3] PyTorch: Profiler Recipe; [4] PyTorch: CUDA semantics, memory management and PYTORCH_CUDA_ALLOC_CONF (read 4 August 2026); [5] PyTorch 2.5 release: FlexAttention and Flight Recorder (read 4 August 2026)