A practical guide to eliminating data transfer bottlenecks in PyTorch — achieving 1.5x speedup with pinned memory, CUDA streams, and GPU Direct Storage. A practical guide to eliminating data transfer bottlenecks in PyTorch — achieving 1.5x speedup with pinned memory, CUDA streams, and GPU Direct Storage. We assumed the GPU was our bottleneck. We were wrong. While training a transformer model, I noticed something strange in the profiler output: the CPU was spending 77% of its time on cudaMemcpyAsync. Our expensive A100 GPU wasn't compute-bound — it was starving for data. 77% of its time cudaMemcpyAsync starving for data This post covers how we diagnosed the problem, fixed it with three increasingly aggressive optimizations, and hit the next wall. If you're training models on large datasets and haven't profiled your data pipeline, you might be leaving significant performance on the table. The Setup We're training nanoTabPFN, a transformer for tabular data. Training data lives in HDF5 files: 30,000 samples, each with 5,000 rows and 5 features. Hardware: NVIDIA A100-SXM4–80GB. nanoTabPFN The original data loading code was textbook PyTorch: with h5py.File(filename, "r") as f:
    for step in range(num_steps):
        x = torch.from_numpy(f["X"][ptr:end])
        y = torch.from_numpy(f["y"][ptr:end])
        yield dict(x=x.to(device), y=y.to(device)) with h5py.File(filename, "r") as f:
    for step in range(num_steps):
        x = torch.from_numpy(f["X"][ptr:end])
        y = torch.from_numpy(f["y"][ptr:end])
        yield dict(x=x.to(device), y=y.to(device)) Simple. Correct. And devastatingly slow. Profile First, Optimize Later Before touching any code, we ran PyTorch's built-in profiler: from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
) as prof:
    train(model, prior)

print(prof.key_averages().table(sort_by="cpu_time_total")) from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
) as prof:
    train(model, prior)

print(prof.key_averages().table(sort_by="cpu_time_total")) The results were shocking: Operation

CPU Time

% of Total



cudaMemcpyAsync

44,084ms

76.78%



cudaMalloc

7,081ms

12.33%



cudaLaunchKernel

645ms

1.12%



aten::bmm

180ms

0.31% Operation

CPU Time

% of Total



cudaMemcpyAsync

44,084ms

76.78%



cudaMalloc

7,081ms

12.33%



cudaLaunchKernel

645ms

1.12%



aten::bmm

180ms

0.31% Operation

CPU Time

% of Total Operation Operation CPU Time CPU Time % of Total % of Total cudaMemcpyAsync

44,084ms

76.78% cudaMemcpyAsync cudaMemcpyAsync cudaMemcpyAsync 44,084ms 44,084ms 76.78% 76.78% 76.78% cudaMalloc

7,081ms

12.33% cudaMalloc cudaMalloc cudaMalloc 7,081ms 7,081ms 12.33% 12.33% cudaLaunchKernel

645ms

1.12% cudaLaunchKernel cudaLaunchKernel cudaLaunchKernel 645ms 645ms 1.12% 1.12% aten::bmm

180ms

0.31% aten::bmm aten::bmm aten::bmm 180ms 180ms 0.31% 0.31% The GPU was doing matrix multiplications in milliseconds while the CPU spent 44 seconds copying data. milliseconds 44 seconds Understanding the Problem The .to(device) call in PyTorch is synchronous by default. Here's the hidden pipeline: .to(device) h5py reads from disk → CPU memory (pageable)
PyTorch allocates → CPU staging buffer
cudaMemcpy → GPU memory (blocks until complete)
GPU computes → while CPU waits for step 1 h5py reads from disk → CPU memory (pageable) h5py reads from disk PyTorch allocates → CPU staging buffer PyTorch allocates cudaMemcpy → GPU memory (blocks until complete) cudaMemcpy (blocks until complete) GPU computes → while CPU waits for step 1 GPU computes The GPU sits idle during steps 1–3. With 5,000-row samples at float32, each batch transfer is ~120MB. That's 12GB of sequential transfers over 100 steps. Fix #1: Pinned Memory + Non-blocking Transfers The first optimization: use page-locked (pinned) memory with async transfers. # Before: synchronous, pageable memory
x = torch.from_numpy(x_np).to(device)

# After: pinned memory, async transfer
x = torch.from_numpy(x_np).pin_memory().to(device, non_blocking=True) # Before: synchronous, pageable memory
x = torch.from_numpy(x_np).to(device)

# After: pinned memory, async transfer
x = torch.from_numpy(x_np).pin_memory().to(device, non_blocking=True) Why this works: Pinned memory is DMA-accessible — the GPU can read it directly without CPU intervention. Combined with non_blocking=True, the transfer happens in the background while the CPU continues working. Why this works: non_blocking=True Impact: cudaMemcpyAsync time dropped from 44s to ~4s. Impact: cudaMemcpyAsync Fix #2: CUDA Streams for True Overlap Non-blocking transfers alone aren't enough. By default, operations on the same CUDA stream are serialized. We need a separate stream for data transfer: class PriorDumpDataLoader:
    def __init__(self, ...):
        self.transfer_stream = torch.cuda.Stream()
    
    def __iter__(self):
        # Pre-fill buffer with first batches
        vram_buffer = [self._load_to_vram(f) for _ in range(prefetch)]
        
        for step in range(num_steps):
            batch = vram_buffer.pop(0)  # Already in VRAM
            
            # Prefetch next batch on separate stream
            with torch.cuda.stream(self.transfer_stream):
                next_batch = self._load_to_vram(f)
            vram_buffer.append(next_batch)
            
            # Sync before yielding
            torch.cuda.current_stream().wait_stream(self.transfer_stream)
            yield batch class PriorDumpDataLoader:
    def __init__(self, ...):
        self.transfer_stream = torch.cuda.Stream()
    
    def __iter__(self):
        # Pre-fill buffer with first batches
        vram_buffer = [self._load_to_vram(f) for _ in range(prefetch)]
        
        for step in range(num_steps):
            batch = vram_buffer.pop(0)  # Already in VRAM
            
            # Prefetch next batch on separate stream
            with torch.cuda.stream(self.transfer_stream):
                next_batch = self._load_to_vram(f)
            vram_buffer.append(next_batch)
            
            # Sync before yielding
            torch.cuda.current_stream().wait_stream(self.transfer_stream)
            yield batch This is double buffering: while the GPU processes batch N, the CPU+DMA engine load batch N+1. The GPU never waits. double buffering Fix #3: GPU Direct Storage (GDS) The ultimate optimization: bypass the CPU entirely. NVIDIA's GPUDirect Storage reads directly from NVMe to GPU memory: import kvikio
import cupy as cp

# Allocate GPU buffer
x_gpu = cp.empty((batch_size, seq_len, features), dtype=cp.float32)

# Direct read: NVMe → GPU (no CPU copy)
with kvikio.CuFile("data.bin", "r") as f:
    f.pread(x_gpu, file_offset=offset)

# Zero-copy to PyTorch
x = torch.as_tensor(x_gpu, device="cuda") import kvikio
import cupy as cp

# Allocate GPU buffer
x_gpu = cp.empty((batch_size, seq_len, features), dtype=cp.float32)

# Direct read: NVMe → GPU (no CPU copy)
with kvikio.CuFile("data.bin", "r") as f:
    f.pread(x_gpu, file_offset=offset)

# Zero-copy to PyTorch
x = torch.as_tensor(x_gpu, device="cuda") The catch: GDS requires raw binary files. HDF5 has headers that need CPU parsing. We added automatic conversion on first run: The catch: def convert_h5_to_raw(h5_filename):
    with h5py.File(h5_filename, "r") as f:
        X = f["X"][:].astype(np.float32)
        y = f["y"][:].astype(np.float32)
    X.tofile(f"{base}_X.bin")
    y.tofile(f"{base}_y.bin") def convert_h5_to_raw(h5_filename):
    with h5py.File(h5_filename, "r") as f:
        X = f["X"][:].astype(np.float32)
        y = f["y"][:].astype(np.float32)
    X.tofile(f"{base}_X.bin")
    y.tofile(f"{base}_y.bin") Results Metric

Baseline

Optimized

Speedup



Total time (100 steps)

68.75s

45.30s

1.52x



cudaMemcpyAsync CPU

44,084ms

268ms

164x



Steps/sec

1.5

2.2

1.47x Metric

Baseline

Optimized

Speedup



Total time (100 steps)

68.75s

45.30s

1.52x



cudaMemcpyAsync CPU

44,084ms

268ms

164x



Steps/sec

1.5

2.2

1.47x Metric

Baseline

Optimized

Speedup Metric Metric Baseline Baseline Optimized Optimized Speedup Speedup Total time (100 steps)

68.75s

45.30s

1.52x Total time (100 steps) Total time (100 steps) 68.75s 68.75s 45.30s 45.30s 1.52x 1.52x 1.52x cudaMemcpyAsync CPU

44,084ms

268ms

164x cudaMemcpyAsync CPU cudaMemcpyAsync CPU cudaMemcpyAsync 44,084ms 44,084ms 268ms 268ms 164x 164x 164x Steps/sec

1.5

2.2

1.47x Steps/sec Steps/sec 1.5 1.5 2.2 2.2 1.47x 1.47x 1.47x Memory transfer overhead dropped from 77% to <1% of CPU time. The New Bottleneck With data loading solved, the profile looks completely different: Operation

CPU Time

% of Total



Command Buffer Full

23,450ms

46.91%



cudaLaunchKernel

10,733ms

21.47%



cudaMalloc

5,607ms

11.22% Operation

CPU Time

% of Total



Command Buffer Full

23,450ms

46.91%



cudaLaunchKernel

10,733ms

21.47%



cudaMalloc

5,607ms

11.22% Operation

CPU Time

% of Total Operation Operation CPU Time CPU Time % of Total % of Total Command Buffer Full

23,450ms

46.91% Command Buffer Full Command Buffer Full Command Buffer Full 23,450ms 23,450ms 46.91% 46.91% 46.91% cudaLaunchKernel

10,733ms

21.47% cudaLaunchKernel cudaLaunchKernel cudaLaunchKernel 10,733ms 10,733ms 21.47% 21.47% cudaMalloc

5,607ms

11.22% cudaMalloc cudaMalloc cudaMalloc 5,607ms 5,607ms 11.22% 11.22% The GPU is now saturated. "Command Buffer Full" means the GPU can't keep up with kernel submissions. This is exactly what we want — the GPU is the bottleneck, not data loading. saturated The remaining compute bottleneck is attention (aten::bmm at 45% CUDA time). With 5,000-row sequences, attention's O(n²) scaling dominates. Flash Attention is the next optimization. aten::bmm Key Takeaways Async is not automatic. non_blocking=True does nothing without proper stream management. Async is not automatic. non_blocking=True Pinned memory matters. 10x+ difference for large transfers. Pinned memory matters. GDS has constraints. True zero-copy requires raw binary files, GDS-compatible NVMe, and proper alignment. GDS has constraints. Know when to stop. Once you're GPU-bound, data loading optimizations won't help. Move to model architecture changes. Know when to stop. Quick Reference Technique

What it does

When to use



pin_memory()

Page-locked CPU memory

Always for GPU training



non_blocking=True

Async H2D transfer

With CUDA streams



CUDA Streams

Parallel transfer/compute

Large batch sizes



Double buffering

Prefetch next batch

I/O-bound workloads



GDS (kvikio)

Disk → GPU direct

Large sequential reads Technique

What it does

When to use



pin_memory()

Page-locked CPU memory

Always for GPU training



non_blocking=True

Async H2D transfer

With CUDA streams



CUDA Streams

Parallel transfer/compute

Large batch sizes



Double buffering

Prefetch next batch

I/O-bound workloads



GDS (kvikio)

Disk → GPU direct

Large sequential reads Technique

What it does

When to use Technique Technique What it does What it does When to use When to use pin_memory()

Page-locked CPU memory

Always for GPU training pin_memory() pin_memory() pin_memory() Page-locked CPU memory Page-locked CPU memory Always for GPU training Always for GPU training non_blocking=True

Async H2D transfer

With CUDA streams non_blocking=True non_blocking=True non_blocking=True Async H2D transfer Async H2D transfer With CUDA streams With CUDA streams CUDA Streams

Parallel transfer/compute

Large batch sizes CUDA Streams CUDA Streams Parallel transfer/compute Parallel transfer/compute Large batch sizes Large batch sizes Double buffering

Prefetch next batch

I/O-bound workloads Double buffering Double buffering Prefetch next batch Prefetch next batch I/O-bound workloads I/O-bound workloads GDS (kvikio)

Disk → GPU direct

Large sequential reads GDS (kvikio) GDS (kvikio) Disk → GPU direct Disk → GPU direct Large sequential reads Large sequential reads Code All code is available at github.com/stprnvsh/nanoTabPFN: github.com/stprnvsh/nanoTabPFN # Baseline
python train.py --profile --steps=100 --batch-size=6

# Optimized with GDS
python train_optimized.py --gds-bin --batch-size=4 --steps=200

# With Flash Attention
python train_optimized.py --flash --gds-bin --batch-size=8 --steps=200 # Baseline
python train.py --profile --steps=100 --batch-size=6

# Optimized with GDS
python train_optimized.py --gds-bin --batch-size=4 --steps=200

# With Flash Attention
python train_optimized.py --flash --gds-bin --batch-size=8 --steps=200

Our GPU Was Idle 77% of the Time. Here's How We Fixed It

About Author

Comments

TOPICS

THIS ARTICLE WAS FEATURED IN

Related Stories

5 Ways Artificial Intelligence Is (Quietly) Changing Libraries

Mine Intelligence Instead of Bitcoin: The Future of Green Computing

Why Salesforce Chose London for Its First AI Center and $4 Billion AI Investment

Human Preferences Help Scientists Train AI 30x Faster Than Before

How Gradient-Free Training Could Decentralize AI

Cloudflare Gives Website Owners Option to Charge OpenAI Bots for Scraping

5 Ways Artificial Intelligence Is (Quietly) Changing Libraries

Mine Intelligence Instead of Bitcoin: The Future of Green Computing

Why Salesforce Chose London for Its First AI Center and $4 Billion AI Investment

Human Preferences Help Scientists Train AI 30x Faster Than Before

How Gradient-Free Training Could Decentralize AI

Cloudflare Gives Website Owners Option to Charge OpenAI Bots for Scraping

Light-Mode

Classic

Newspaper

Minty

Dark-Mode

Neon Noir

Minty

HN StartUps