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