AI Infra Debug

Strategy

1. Baseline measurements: “Before changing anything, I measured X, Y, Z. The bottleneck appears to be A.”
2. Concrete improvement: For example: reduced latency, improved GPU utilization, removed a memory bottleneck, improved batching, reduced redundant work, or found a config issue.
3. Final writeup: 1–2 page writeup

Report

Summary:

• Investigated [system/component].
• Found bottleneck in [area].
• Implemented/tested [change].
• Result: [metric before] -> [metric after].
• Remaining opportunities: [ranked list].

Method:

• Baseline setup
• Metrics collected
• Tools used
• Assumptions

Findings:

1. ...
2. ...
3. ...

Changes made:

• ...

Risks / tradeoffs:

• ...

Recommended next steps:

1. ...
2. ...
3. ...

Strategy

Identify Metrics:

• What matters most:
  • latency
  • throughput
  • GPU utilization
  • memory
  • cost
  • reliability
  • developer simplicity
• What is the current behavior, and what would count as a useful improvement

Useful Commands:

• watch -n 1 nvidia-smi: Shows GPU memory, utilization, power, running processes.
• htop: Shows CPU usage.
• iostat -x 1: Shows disk bottlenecks, if available.
• df -h: Checks disk space.
• du -sh .: Checks directory size.

For Python timing

import time

t0 = time.perf_counter()
# work here
t1 = time.perf_counter()

print(f"Elapsed: {t1 - t0:.3f}s")

For GPU timing, use CUDA synchronization

Without torch.cuda.synchronize(), GPU timing can lie because CUDA work is asynchronous.

import time
import torch

torch.cuda.synchronize()
t0 = time.perf_counter()

# GPU work here

torch.cuda.synchronize()
t1 = time.perf_counter()

print(f"GPU elapsed: {t1 - t0:.3f}s")

A model inference endpoint is slow

This means: an API endpoint receives input, runs a model, returns output, and response time is bad.

First, break one request into stages:

• request receive
• input validation
• preprocessing
• CPU -> GPU transfer
• model inference
• GPU -> CPU transfer
• postprocessing
• response serialization
• network response

Add timers around each stage. Example:

import time
import torch

def now():
    return time.perf_counter()

def predict(request):
    t0 = now()

    # parse input
    x = parse_request(request)
    t1 = now()

    # preprocess
    x = preprocess(x)
    t2 = now()

    # move to GPU
    x = x.to("cuda")
    t3 = now()

    # model inference
    torch.cuda.synchronize()
    t4 = now()

    with torch.inference_mode():
        y = model(x)

    torch.cuda.synchronize()
    t5 = now()

    # postprocess
    result = postprocess(y)
    t6 = now()

    print({ "parse": t1 - t0, "preprocess": t2 - t1, "to_gpu": t3 - t2, "model": t5 - t4, "postprocess": t6 - t5, "total": t6 - t0, })

    return result

Common causes:

• The model is loaded on every request.
• The model is not on GPU.
• The input is moved CPU -> GPU repeatedly in a bad way.
• The endpoint processes one item at a time instead of batching.
• The code uses training mode instead of inference mode.
• Preprocessing is slow.
• Postprocessing is slow.
• Serialization is slow.
• The first request is slow because of warmup.

Make sure model is loaded once, not per request:

model = load_model()
model.to("cuda")
model.eval()

Use inference mode:

with torch.inference_mode():
    output = model(input)

Warm up the model:

with torch.inference_mode():
    for _ in range(3):
        _ = model(dummy_input)

Avoid per-request heavyweight setup.

Bad:

def predict(x):
    model = load_model()
    model.to("cuda")
    return model(x)

Good:

model = load_model().to("cuda")
model.eval()

def predict(x):
    with torch.inference_mode():
        return model(x)

Metrics:

• Baseline p50 latency: X ms
• Baseline p95 (95% of data falls at or below) latency: Y ms

Breakdown:

• preprocessing: A ms
• model inference: B ms
• postprocessing: C ms

A batch job uses too much GPU memory

A script runs a model over many inputs and either crashes with CUDA OOM or uses more memory than expected.

Add watch:

watch -n 1 nvidia-smi

Measure peak GPU memory

import torch

torch.cuda.reset_peak_memory_stats()

# run workload

print("allocated GB:", torch.cuda.memory_allocated() / 1e9)
print("reserved GB:", torch.cuda.memory_reserved() / 1e9)
print("peak allocated GB:", torch.cuda.max_memory_allocated() / 1e9)

Print summary:

print(torch.cuda.memory_summary())

Common causes:

• Batch size is too large.
• Code tracks gradients during inference.
• The model is in training mode.
• Outputs are stored without detaching from computation graph.
• A list keeps accumulating GPU tensors.
• Data is copied unnecessarily.
• Multiple model copies exist on GPU.
• Intermediate tensors are not released.
• The job uses float32 when float16/bfloat16 is enough.
• First safe fixes

Common approaches:

For inference:

model.eval()

with torch.inference_mode():
    output = model(input)

If they are using:

with torch.no_grad():

That is also okay, but inference_mode() is usually stronger for inference.

Reduce batch size:

batch_size = batch_size // 2

Use mixed precision if appropriate:

with torch.inference_mode():
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        output = model(input)

Or on newer NVIDIA GPUs, sometimes:

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    output = model(input)

Do not store GPU tensors forever.

Bad:

outputs = []

for batch in loader:
    output = model(batch)
    outputs.append(output)

Better:

outputs = []

for batch in loader:
    with torch.inference_mode():
        output = model(batch)

    outputs.append(output.detach().cpu())

If you only need final values, move them to CPU or write them to disk.

Important warning, do not treat torch.cuda.empty_cache() as a real fix. It can help release cached memory back to CUDA, but it usually does not solve the real memory bug.

Metrics:

Baseline peak GPU memory: X GB
OOM occurs at batch size: N

Findings:
- gradients were being tracked during inference
- outputs were stored on GPU
- batch size was too large
- memory grows each iteration, suggesting accumulation

Change:
- added inference_mode
- moved outputs to CPU
- reduced batch size
- tested autocast

Result:
- peak memory reduced from X GB to Y GB
- job completed successfully

A data pipeline is bottlenecked

The GPU is waiting because data loading, image decoding, preprocessing, network fetch, or disk reads are too slow.

How you recognize it

watch -n 1 nvidia-smi

If GPU utilization jumps between 0% and 60% or stays low, but CPU is busy, the GPU may be starving.

Then time the stages:

for batch in loader:
    t0 = time.perf_counter()

    # batch already loaded here
    t1 = time.perf_counter()

    batch = batch.to("cuda")
    t2 = time.perf_counter()

    with torch.inference_mode():
        output = model(batch)

    torch.cuda.synchronize()
    t3 = time.perf_counter()

    print({
        "to_gpu": t2 - t1,
        "model": t3 - t2,
    })

Measure loading time:

data_iter = iter(loader)

for _ in range(10):
    t0 = time.perf_counter()
    batch = next(data_iter)
    t1 = time.perf_counter()

    print("load batch:", t1 - t0)

Common causes

• num_workers is too low.
• Data is read from slow remote storage.
• Images/videos are decoded on CPU too slowly.
• Transforms are expensive.
• Batch size is too small.
• Data is not prefetched.
• Data is copied CPU -> GPU inefficiently.
• Dataset has many tiny files.
• Safe fixes

If using PyTorch DataLoader, inspect:

DataLoader(
    dataset,
    batch_size=...,
    num_workers=...,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=2,
)

Try increasing workers (too many workers can hurt too, test 4 then 8):

num_workers=4 or 8

Use pinned memory for GPU transfer:

loader = DataLoader(dataset, pin_memory=True)

Then:

batch = batch.to("cuda", non_blocking=True)

Cache expensive preprocessing if repeated.

Bad:

for epoch in range(10):
    image = decode_and_resize(path)

Better:

decode/resize once
save preprocessed form
reuse it

If there are many tiny files they may cause metadata/read overhead. A packed format or local cache could improve throughput.

Baseline:

• average data loading time: X ms/batch
• average model time: Y ms/batch
• GPU utilization: low / bursty

Finding:

• GPU is waiting on data loading/preprocessing.

Change tested:

• increased DataLoader workers from A to B
• enabled pin_memory
• enabled persistent_workers
• tested local cache

Result:

• batch loading improved from X to Y
• GPU utilization improved from A% to B%
• throughput improved from X samples/s to Y samples/s

A PyTorch script has poor throughput

The script works, but samples/sec is too low.

What you measure

• Throughput:

import time
import torch

num_samples = 0

torch.cuda.synchronize()
t0 = time.perf_counter()

for batch in loader:
    batch = batch.to("cuda")

    with torch.inference_mode():
        output = model(batch)

    num_samples += len(batch)

torch.cuda.synchronize()
t1 = time.perf_counter()

print("samples/sec:", num_samples / (t1 - t0))

Common causes

• Batch size too small.
• Running one sample at a time.
• Python loops around tensor operations.
• Using gradients during inference.
• Calling .cpu(), .numpy(), or .item() inside the hot loop.
• Excessive logging.
• DataLoader bottleneck.
• Model not on GPU.
• Input not on GPU.
• Repeated model initialization.
• No warmup.

Things to search for in code that if in hot path are suspicious.

.item()
.cpu()
.numpy()
print(...)
for sample in batch:
model = ...
load_model(...)

Bad:

for sample in samples:
    output = model(sample.to("cuda"))

Better:

batch = torch.stack(samples).to("cuda")

with torch.inference_mode():
    output = model(batch)

Bad:

for x in tensor:
    result.append(slow_python_function(x))

Better: use tensor operations if possible.

Safe PyTorch approaches

• Inference mode:

model.eval()

with torch.inference_mode():
    output = model(input)

Batching:

# Instead of N single inferences
# Do one batched inference
output = model(batch)

Mixed precision, if acceptable:

with torch.inference_mode():
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        output = model(input)

Avoid synchronizing too often. These can force GPU sync(Sometimes okay, but bad inside tight loops):

loss.item()
tensor.cpu()
tensor.numpy()
print(tensor)

What you report

• Baseline throughput: X samples/sec

Findings:

• script was doing per-sample inference
• hot loop included .cpu() / .item()
• gradients were enabled during inference
• DataLoader was slower than model execution

Changes:

• batched inference
• added inference_mode
• removed unnecessary CPU sync from hot loop
• adjusted batch size

Result:

• throughput improved from X to Y samples/sec

A service has bad latency under load

One request may be okay, but many concurrent users make it slow.

What to measure: latency under concurrency:

• single request latency
• p50 latency
• p95 latency
• p99 latency
• requests per second
• error rate

Use custom tools, wrk, hey, ab, or Locust.

Example with hey.wrk:

hey -n 1000 -c 20 http://localhost:8000/predict
wrk -t4 -c32 -d60s http://localhost:8000/predict

Common causes

• Server handles requests serially.
• Model inference blocks the event loop.
• Too many concurrent requests overwhelm GPU memory.
• No batching.
• No queue/backpressure.
• Too many workers each load their own model copy.
• Database/object storage call is slow.
• Cold starts.
• Large response serialization.

If it is FastAPI/async Python, look for blocking work inside async endpoint.

Suspicious:

@app.post("/predict")
async def predict(req):
    result = slow_blocking_model_call(req)
    return result

• Async does not magically make CPU/GPU work non-blocking.
• Check number of workers. If each worker loads a model, this can duplicate GPU memory:

ps aux | grep python
nvidia-smi

• If you see many Python processes each using GPU memory, that may be the issue.

Safe fixes

• Add simple concurrency control to reduce overload and improve p95/p99, even if p50 stays similar.:

import asyncio

semaphore = asyncio.Semaphore(4)

@app.post("/predict")
async def predict(req):
    async with semaphore:
        return run_model(req)

• Add batching if system supports it.
• Add queue/backpressure: If too many requests arrive, queue them or reject gracefully instead of letting latency explode.
• Separate model time from request overhead.

What to report

• Baseline under 20 concurrent clients:
• throughput: X req/s
• p50: A ms
• p95: B ms
• errors: C%

Finding:

• latency explodes after concurrency N
• GPU memory saturates / CPU saturates / requests serialize / event loop blocks

Change:

• added concurrency limit
• reduced worker count
• tested batching
• moved blocking work out of async path

Result:

• p95 improved from X to Y
• error rate dropped from A% to B%