GPU Memory Optimization

Work carried out at SEGULA Technologies with the CNRS Hubert Curien Laboratory, first as an intern and then as an R&D engineer, on GPU memory management during neural network training.

The practical question: when training is limited by GPU memory rather than by compute, what can be changed in how memory is allocated and moved to lift that limit?

Training deep networks is frequently bounded by GPU memory rather than by arithmetic throughput. Activations stored for the backward pass dominate the footprint, and that footprint scales with batch size, depth and input resolution.

When memory is the binding constraint, the usual responses — reduce batch size, reduce resolution, use a smaller model — all degrade the result. Understanding where the memory actually goes is the prerequisite for any better option.

• Activation storage dominates training memory
• Allocation patterns are opaque without instrumentation
• Fragmentation reduces usable memory below the nominal capacity
• Host-to-device transfers can stall the GPU when poorly scheduled

• Instrument and visualise memory flows during training
• Identify where allocation patterns waste capacity
• Develop memory-management algorithms that reduce peak usage
• Implement them at the CUDA and cuDNN level rather than the framework level
• Use CPU multithreading to overlap transfers with computation

• Fixed physical GPU memory on the available hardware
• Changes must not alter training results — memory management is not allowed to change numerics
• Instrumentation overhead must be small enough not to distort what it measures
• Implementation must work with cuDNN's own workspace requirements

1. 01

   Training Process

   Forward and backward passes

2. 02

   Instrumentation Layer

   Records allocations, frees and transfers

3. 03

   Memory Flow Analysis

   Lifetime, peak usage and fragmentation

4. 04

   Allocation Strategy

   Reuse, scheduling and transfer overlap

5. 05

   Visualization Tools

   Timeline views of memory over training steps

• Allocation and free events captured with timestamps during training
• Tensor lifetimes reconstructed from the event stream
• Peak-usage points located and attributed to specific training phases
• Allocation strategies evaluated against the recorded workload
• Transfers scheduled on separate streams to overlap with computation

• Memory pooling to reduce allocator overhead and fragmentation
• Reuse of buffers whose lifetimes do not overlap
• Asynchronous host-to-device transfers overlapped with kernel execution
• CPU multithreading for data preparation so the GPU is not starved
• cuDNN workspace sizing chosen deliberately rather than left at default

• Delivered as instrumentation and memory-management components used within training workflows
• Visualization tools used to diagnose memory behaviour on specific training runs
• Industrial R&D context; not released publicly

• Numerical equivalence verified: optimized runs must produce the same results as unoptimized ones
• Instrumentation overhead measured and reported alongside results
• Out-of-memory conditions handled with diagnostic output rather than an opaque crash

Memory work is only credible when measured on the same hardware, model, batch size and input resolution, and when the instrumentation cost is stated. Comparisons hold all of those constant and report peak memory and step time together, since improving one at the cost of the other is not an improvement.

Measurements belong to the industrial R&D context in which the work was carried out and are not published here.

Benchmarking in progress. Results will be published with the hardware, model, runtime and measurement conditions that produced them.

• Memory behaviour cannot be reasoned about from source code alone; it has to be observed.
• A timeline visualisation of allocations makes bottlenecks obvious that are invisible in aggregate statistics.
• Peak memory and step time must be optimised together; improving one in isolation usually moves the cost rather than removing it.
• Verifying numerical equivalence after every change is what makes low-level memory work safe to do at all.

• Extension to multi-GPU training where transfers between devices dominate
• Automated selection of allocation strategy from an observed workload profile
• Integration of the visualisation tooling with standard framework profilers