Explain how autograd works. What do requires_grad, detach, no_grad and retain_graph do?

PyTorch builds a dynamic computation graph during the forward pass, recording each operation and its inputs. Calling .backward() traverses it in reverse applying the chain rule and accumulating into .grad. requires_grad marks a tensor as needing gradients; detach() returns a tensor sharing storage but cut from the graph; no_grad() disables graph construction entirely for a block; retain_graph=True keeps the graph alive so you can backward through it twice. Construct | Effect | Typical use. requires_grad=True | Track ops on this tensor, allocate .grad | Model parameters, adversarial inputs. .detach() | Same storage, cut from the graph | Logging, stopping gradient through a branch. torch.no_grad() | No graph built at all inside the block | Inference, validation, manual param updates. torch.inference_mode() | Stronger than no_grad; also disables version counters | Production inference (slightly faster). retain_graph=True | Keep buffers after backward | Multiple backwards through one graph (GANs, meta-learning). create_graph=True | Make the backward itself differentiable | Second-order gradients, MAML. .item() | Extract a Python scalar, breaking the graph | Logging a loss without holding the graph The memory-leak classic: total_loss += loss inside a training loop keeps every step's computation graph alive, because loss still references it. Memory grows until OOM. The fix is total_loss += loss.item() (or .detach()). This is one of the most common real PyTorch bugs and a favourite interview question. Why 'dynamic graph' matters: The graph is built fresh on every forward pass, which is why control flow works naturally - an if or a while in your forward just changes which operations get recorded. That is the core difference from static-graph frameworks, and it is why PyTorch

Write a complete, correct PyTorch training loop from memory.

Set the model to train mode, iterate batches, move to device, zero gradients, forward, compute loss, backward, clip gradients, optimiser step, scheduler step. Then a validation pass under no_grad() in eval mode, tracking the best checkpoint. The details that separate correct from nearly-correct are zero_grad(set_to_none=True), .item() for accumulation, model.train()/eval() toggling, and stepping the scheduler at the right granularity. Detail | Why it matters. model.train() / model.eval() | Toggles dropout and BatchNorm behaviour - the classic silent bug. zero_grad(set_to_none=True) | Frees memory and skips a kernel vs. filling with zeros. .item() when accumulating | Otherwise every step's graph stays alive → OOM. non_blocking=True + pin_memory | Overlaps host-to-device copy with compute. clip_grad_norm_ before step() | After backward, before the update - order matters. Scheduler granularity | Per-step for warmup/OneCycle; per-epoch for StepLR. Save the optimiser state too | Required to resume training correctly The DataLoader settings that actually matter: num_workers > 0 (usually 4–8) so data loading overlaps with GPU compute; pin_memory=True for faster host-to-device transfer; persistent_workers=True to avoid re-spawning every epoch; drop_last=True for training when using BatchNorm so a final batch of size 1 does not blow up. A GPU sitting at 30% utilisation is almost always a data-loading problem, not a model problem. The sanity check to mention: “Before a real run I overfit a single batch - if the model cannot drive the loss to near zero on 8 examples, there is a bug and no amount of tuning will fix it.” This takes 30 seconds and catches wrong loss functions, wrong label alignment, frozen parameters and broken gradients. Volunteering it signals real experience.

A model's loss is NaN after 50 steps and another isn't learning at all. Debug both.

For NaN: check for division by zero or log(0) in the loss, exploding gradients (log the gradient norm), fp16 overflow, and bad input data. For not-learning: verify the loss is actually connected to the parameters, that zero_grad and step are both called, that the learning rate is sane, that labels align with predictions, and that the model can overfit a single batch. Symptom | Likely cause | Check. NaN loss immediately | log(0), division by zero, or bad input data | torch.isnan(x).any(); add ε to denominators. NaN after N steps | Exploding gradients or fp16 overflow | Log grad norm each step; switch to bf16; clip. Loss flat at ln(k) | Gradients are not reaching the parameters | Check for .detach(), missing step(), requires_grad=False. Loss decreases then plateaus high | LR too high or too low; capacity too small | LR range test; try to overfit one batch. Train loss down, val loss up | Overfitting | Regularise, augment, early stop. Val loss lower than train | Dropout active in train only, or a split bug | Normal with heavy dropout; check for leakage. Loss oscillates wildly | LR too high or batch too small | Lower LR; increase batch or accumulate. Accuracy stuck at 1/k | Label/prediction misalignment | Print predictions vs labels for one batch The 'expected initial loss' trick: For balanced k-class cross-entropy, an untrained model should output roughly uniform probabilities, giving a loss of ln(k) - 2.30 for 10 classes, 6.91 for 1000. If your initial loss is 0.02 you have a label leak; if it is 15 your initialisation or loss is wrong. This one check takes seconds and catches an enormous class of bugs. Very few candidates know it. Say the process, not the guess: “I'd check inputs, then shapes, then initial loss against ln(k), then gradient

Write a custom Dataset and DataLoader for variable-length text, including a collate function.

Implement __len__ and __getitem__ returning a single unbatched example, then write a collate_fn that pads sequences to the batch maximum and builds an attention mask. Pad to the longest in the batch rather than a global maximum, and use a length-grouped sampler so batches contain similar lengths - that alone often cuts wasted compute by 30–50%. Why length grouping matters so much: With random batching, one 500-token document forces all 31 other sequences in the batch to be padded to 500. If the median is 87 tokens, you are spending most of your compute on padding. Sorting within a large shuffled 'mega-batch' preserves randomness across epochs while making each batch homogeneous. This is the standard trick in every serious NLP training pipeline and it is a strong thing to bring up unprompted. DataLoader argument | What it does | Recommendation. num_workers | Parallel data loading processes | 4–8; 0 for debugging (real tracebacks). pin_memory | Page-locked host memory for fast H2D copy | True with CUDA. persistent_workers | Keep workers alive across epochs | True when num_workers > 0. prefetch_factor | Batches prefetched per worker | 2–4. drop_last | Discard the final partial batch | True for training with BatchNorm. collate_fn | Assemble a batch from samples | Required for variable-length data Worker gotchas: Each worker is a separate process, so each gets its own copy of the dataset object - holding a large in-memory list means N× the RAM. Store paths or use a memory-mapped format instead. Since PyTorch 1.9 each worker reseeds random, torch and NumPy's global RNG, but an RNG object built in the dataset's __init__ (say np.random.default_rng(0)) is copied into every worker in the same state - so every worker generates

Your training epoch takes 40 minutes and GPU utilisation is 30%. Find and fix the bottleneck.

30% utilisation almost always means the GPU is starved by the input pipeline, not that the model is slow. Check data loading first (num_workers, pin_memory, expensive per-item transforms, slow storage), then synchronisation points (.item(), .cpu(), printing inside the loop), then the model itself. Use the PyTorch profiler to see the actual gaps rather than guessing. Cause | Symptom | Fix. num_workers=0 | GPU idle between batches | Set to 4–8 with persistent_workers. Expensive per-item transform | High worker CPU, GPU idle | Precompute; move augmentation to GPU. Slow storage / small files | Workers I/O-bound | WebDataset/tar shards; local NVMe cache. .item() in the loop | Frequent CPU-GPU syncs | Accumulate on GPU; sync once per epoch. print / logging per step | Implicit sync | Log every N steps. Small batch size | GPU underutilised per kernel | Increase batch; use accumulation for memory. fp32 on tensor-core hardware | Half the achievable throughput | Enable bf16 autocast. Unfused elementwise ops | Many small kernels | torch.compile The mental model: GPU utilisation tells you who is waiting. Low utilisation = the GPU is starved → fix the input pipeline. High utilisation but slow = the model or kernels are the bottleneck → mixed precision, compile, better kernels, or a smaller model. Stating that dichotomy first shows you will not waste a week optimising the wrong layer. Beware hidden synchronisation: .item(), .cpu(), .numpy(), print(tensor), if loss > x: and even len() on some GPU structures force the CPU to wait for the GPU to finish. One of these inside the inner loop can halve throughput and it is invisible in the code. Accumulate metrics on-device and synchronise once per epoch.

PyTorch Interview Questions: Autograd, Training Loops & Debugging

AdvancedautogradpytorchgradientsAsked at Meta AI, NVIDIA

30-second answer

PyTorch builds a dynamic computation graph during the forward pass, recording each operation and its inputs. Calling .backward() traverses it in reverse applying the chain rule and accumulating into .grad. requires_grad marks a tensor as needing gradients; detach() returns a tensor sharing storage but cut from the graph; no_grad() disables graph construction entirely for a block; retain_graph=True keeps the graph alive so you can backward through it twice.

pythonInspecting the graph autograd builds

import torch

x = torch.tensor([2.0], requires_grad=True)
y = x ** 2
z = 3 * y + 1

print(f"y.grad_fn = {y.grad_fn}")            # <PowBackward0>
print(f"z.grad_fn = {z.grad_fn}")            # <AddBackward0>
print(f"parents   = {z.grad_fn.next_functions}")
# ((<MulBackward0>, 0), (None, 0))
# The graph is: x -> Pow -> Mul -> Add -> z, built during the FORWARD pass.

z.backward()
print(f"dz/dx = {x.grad}")                   # 6x = 12

# Gradients ACCUMULATE -- this is why zero_grad() exists
z2 = 3 * x ** 2 + 1
z2.backward()
print(f"after second backward: {x.grad}")    # 24, not 12!

Construct	Effect	Typical use
`requires_grad=True`	Track ops on this tensor, allocate `.grad`	Model parameters, adversarial inputs
`.detach()`	Same storage, cut from the graph	Logging, stopping gradient through a branch
`torch.no_grad()`	No graph built at all inside the block	Inference, validation, manual param updates
`torch.inference_mode()`	Stronger than no_grad; also disables version counters	Production inference (slightly faster)
`retain_graph=True`	Keep buffers after backward	Multiple backwards through one graph (GANs, meta-learning)
`create_graph=True`	Make the backward itself differentiable	Second-order gradients, MAML
`.item()`	Extract a Python scalar, breaking the graph	Logging a loss without holding the graph

pythondetach vs no_grad - a distinction people get wrong

import torch

x = torch.tensor([2.0], requires_grad=True)

# detach(): the OUTPUT is cut from the graph, but x still tracks gradients
d = x.detach()
print(d.requires_grad)                       # False
d2 = d * 3
print(d2.requires_grad)                      # False -- nothing downstream tracks

# no_grad(): NOTHING inside the block is tracked, regardless of inputs
with torch.no_grad():
    n = x * 3
print(n.requires_grad)                       # False

# The key difference: detach shares STORAGE with the original
d[0] = 99.0
print(x)                                     # tensor([99.], requires_grad=True)
# Mutating a detached tensor mutates the original -- use .clone().detach()
# when you want an independent copy.

# Practical: stopping gradient through one branch (as in a target network)
with torch.no_grad():
    target = target_net(next_state).max(1).values
loss = F.mse_loss(q_net(state).gather(1, action).squeeze(1), reward + gamma * target)
# The target network receives no gradient -- essential for DQN stability.

Say these points

The graph is built during the forward pass and traversed in reverse by .backward().
Gradients accumulate - hence zero_grad() every step.
detach() cuts the output from the graph but shares storage; no_grad() disables tracking wholesale.
total_loss += loss without .item() leaks the entire graph each step.
Dynamic graphs are what make Python control flow work inside forward.

Avoid these mistakes

Accumulating loss tensors instead of floats.
Assuming detach() gives an independent copy - it shares storage.
Forgetting zero_grad() and silently summing gradients across steps.

Expect these follow-ups

Why is gradient accumulation across micro-batches useful, and what must you divide by?
How would you compute a second-order gradient?
What is a graph break in torch.compile and why does it hurt?

Intermediatepytorchtraining-loopbest-practicesAsked at Meta AI, Tesla

30-second answer

Set the model to train mode, iterate batches, move to device, zero gradients, forward, compute loss, backward, clip gradients, optimiser step, scheduler step. Then a validation pass under no_grad() in eval mode, tracking the best checkpoint. The details that separate correct from nearly-correct are zero_grad(set_to_none=True), .item() for accumulation, model.train()/eval() toggling, and stepping the scheduler at the right granularity.

pythonA training loop with every detail correct

import torch, torch.nn as nn
from torch.amp import autocast

def train_one_epoch(model, loader, criterion, optimizer, scheduler,
                    device, max_grad_norm=1.0, accum_steps=1):
    model.train()                                       # dropout ON, BN uses batch stats
    total_loss, n_seen = 0.0, 0
    optimizer.zero_grad(set_to_none=True)               # set_to_none is faster than =0

    for step, (x, y) in enumerate(loader):
        x = x.to(device, non_blocking=True)             # non_blocking + pin_memory
        y = y.to(device, non_blocking=True)             # overlaps H2D with compute

        with autocast(device_type=device.type, dtype=torch.bfloat16):
            out = model(x)
            loss = criterion(out, y)

        (loss / accum_steps).backward()                 # divide for accumulation

        if (step + 1) % accum_steps == 0:
            nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)
            if scheduler is not None:
                scheduler.step()                        # per-STEP for warmup schedules

        total_loss += loss.item() * len(y)              # .item() -- do NOT keep the graph
        n_seen += len(y)

    return total_loss / n_seen


@torch.no_grad()                                        # decorator form is cleaner
def evaluate(model, loader, criterion, device):
    model.eval()                                        # dropout OFF, BN uses running stats
    total_loss, correct, n = 0.0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        with autocast(device_type=device.type, dtype=torch.bfloat16):
            out = model(x)
            loss = criterion(out, y)
        total_loss += loss.item() * len(y)
        correct += (out.argmax(1) == y).sum().item()
        n += len(y)
    return total_loss / n, correct / n


def fit(model, train_loader, val_loader, epochs=20, lr=3e-4, patience=5):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.OneCycleLR(
        optimizer, max_lr=lr, total_steps=epochs * len(train_loader))

    best_val, bad_epochs = float("inf"), 0
    for epoch in range(epochs):
        tr = train_one_epoch(model, train_loader, criterion, optimizer,
                             scheduler, device)
        vl, acc = evaluate(model, val_loader, criterion, device)
        print(f"epoch {epoch:3d}  train {tr:.4f}  val {vl:.4f}  acc {acc:.4f}  "
              f"lr {scheduler.get_last_lr()[0]:.2e}")

        if vl < best_val - 1e-4:
            best_val, bad_epochs = vl, 0
            torch.save({"model": model.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "epoch": epoch, "val_loss": vl}, "best.pt")
        else:
            bad_epochs += 1
            if bad_epochs >= patience:
                print(f"early stopping at epoch {epoch}")
                break

    model.load_state_dict(torch.load("best.pt")["model"])
    return model

Detail	Why it matters
`model.train()` / `model.eval()`	Toggles dropout and BatchNorm behaviour - the classic silent bug
`zero_grad(set_to_none=True)`	Frees memory and skips a kernel vs. filling with zeros
`.item()` when accumulating	Otherwise every step's graph stays alive → OOM
`non_blocking=True` + `pin_memory`	Overlaps host-to-device copy with compute
`clip_grad_norm_` before `step()`	After backward, before the update - order matters
Scheduler granularity	Per-step for warmup/OneCycle; per-epoch for StepLR
Save the optimiser state too	Required to resume training correctly

Say these points

train()/eval() toggling, zero_grad(set_to_none=True), .item() accumulation.
Clip gradients after backward and before optimiser step.
Step per-batch schedulers inside the loop, per-epoch ones outside.
Save optimiser and scheduler state alongside the model to resume correctly.
Overfit a single batch first as a sanity check.

Avoid these mistakes

Forgetting model.eval() in validation, so dropout is active.
Accumulating the loss tensor instead of a float.
Stepping a per-step scheduler once per epoch (or vice versa).

Expect these follow-ups

How would you add distributed data parallel to this loop?
How do you make this resumable from a mid-epoch checkpoint?
Your GPU utilisation is 30%. Where do you look?

AdvanceddebuggingpytorchtrainingAsked at Meta AI, OpenAI

30-second answer

For NaN: check for division by zero or log(0) in the loss, exploding gradients (log the gradient norm), fp16 overflow, and bad input data. For not-learning: verify the loss is actually connected to the parameters, that zero_grad and step are both called, that the learning rate is sane, that labels align with predictions, and that the model can overfit a single batch.

pythonA systematic debugging checklist you can run in order

import torch, math

def diagnose(model, loader, criterion, optimizer, device):
    x, y = next(iter(loader))
    x, y = x.to(device), y.to(device)

    # 1. INPUT SANITY
    print(f"x: shape={tuple(x.shape)} dtype={x.dtype} "
          f"range=[{x.min():.3f},{x.max():.3f}] nan={torch.isnan(x).any().item()}")
    print(f"y: shape={tuple(y.shape)} dtype={y.dtype} "
          f"unique={y.unique().tolist()[:10]}")
    # Common bugs: unnormalised inputs (0-255), labels 1-indexed, wrong dtype

    # 2. FORWARD SANITY
    model.eval()
    with torch.no_grad():
        out = model(x)
    print(f"out: shape={tuple(out.shape)} range=[{out.min():.3f},{out.max():.3f}]")
    # Shape mismatch with y is the single most common cause of silent nonsense

    # 3. INITIAL LOSS -- is it what theory predicts?
    loss = criterion(out, y)
    n_classes = out.shape[-1]
    print(f"initial loss={loss.item():.4f}  expected~{math.log(n_classes):.4f}")
    # For balanced k-class cross-entropy the untrained loss should be ~ln(k).
    # Far off means wrong loss, wrong label encoding, or a bad init.

    # 4. GRADIENTS -- do they exist and are they sane?
    model.train()
    optimizer.zero_grad()
    criterion(model(x), y).backward()
    no_grad, zero_grad_params, big = [], [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:      continue
        if p.grad is None:           no_grad.append(name); continue
        g = p.grad.norm().item()
        if g == 0:                   zero_grad_params.append(name)
        if g > 100 or math.isnan(g): big.append((name, g))
    print(f"params with NO grad: {no_grad[:5]}")        # detached / unused branch
    print(f"params with ZERO grad: {zero_grad_params[:5]}")  # dead ReLU / frozen
    print(f"params with HUGE/NaN grad: {big[:5]}")

    # 5. CAN IT OVERFIT ONE BATCH? The decisive test.
    for i in range(200):
        optimizer.zero_grad()
        l = criterion(model(x), y)
        l.backward(); optimizer.step()
        if i % 50 == 0:
            print(f"  overfit step {i:3d}: loss={l.item():.6f}")
    print("PASS -- the pipeline is wired correctly"
          if l.item() < 0.01 else
          "FAIL -- there is a bug; tuning will not help")

Symptom	Likely cause	Check
NaN loss immediately	log(0), division by zero, or bad input data	`torch.isnan(x).any()`; add ε to denominators
NaN after N steps	Exploding gradients or fp16 overflow	Log grad norm each step; switch to bf16; clip
Loss flat at ln(k)	Gradients are not reaching the parameters	Check for `.detach()`, missing `step()`, `requires_grad=False`
Loss decreases then plateaus high	LR too high or too low; capacity too small	LR range test; try to overfit one batch
Train loss down, val loss up	Overfitting	Regularise, augment, early stop
Val loss lower than train	Dropout active in train only, or a split bug	Normal with heavy dropout; check for leakage
Loss oscillates wildly	LR too high or batch too small	Lower LR; increase batch or accumulate
Accuracy stuck at 1/k	Label/prediction misalignment	Print predictions vs labels for one batch

pythonCatching a NaN the moment it appears

import math, torch

# Development only -- pinpoints the exact op that produced the NaN
torch.autograd.set_detect_anomaly(True)     # ~3x slower; never leave on

# Or hook every module's output
def add_nan_hooks(model):
    def hook(module, inp, out):
        t = out[0] if isinstance(out, tuple) else out
        if torch.is_tensor(t) and not torch.isfinite(t).all():
            raise RuntimeError(f"non-finite output from {module.__class__.__name__}")
    for m in model.modules():
        m.register_forward_hook(hook)

# Log gradient norms every step -- NaNs are almost always preceded by growth
def grad_norm(model):
    return math.sqrt(sum(p.grad.pow(2).sum().item()
                         for p in model.parameters() if p.grad is not None))

# step 40: grad_norm=  1.83
# step 45: grad_norm=  8.91
# step 48: grad_norm= 61.40      <- the divergence started HERE, not at step 50
# step 50: grad_norm=   nan

Say these points

Check inputs → shapes → initial loss vs ln(k) → gradient norms → overfit one batch.
NaN after N steps means growing gradients; log the norm to find the real start.
Params with grad is None indicate a detached or unused branch.
Loss stuck at exactly ln(k) means gradients are not reaching the parameters.
set_detect_anomaly(True) pinpoints the op, at ~3× slowdown.

Avoid these mistakes

Lowering the learning rate reflexively without diagnosing.
Not checking whether labels and predictions are aligned.
Leaving anomaly detection on in a real run.

Expect these follow-ups

Your GPU utilisation is 25% - walk me through diagnosing that.
How would you debug a model that works in eager mode but breaks under torch.compile?
How do you debug a divergence that only occurs in distributed training?

Intermediatedataloaderpytorchdata-pipelineAsked at Hugging Face, Meta AI

30-second answer

Implement __len__ and __getitem__ returning a single unbatched example, then write a collate_fn that pads sequences to the batch maximum and builds an attention mask. Pad to the longest in the batch rather than a global maximum, and use a length-grouped sampler so batches contain similar lengths - that alone often cuts wasted compute by 30–50%.

pythonDataset + collate + length-grouped batching

import torch
from torch.utils.data import Dataset, DataLoader, Sampler
import random

class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length=512):
        self.texts, self.labels = texts, labels
        self.tokenizer, self.max_length = tokenizer, max_length
        # Cache lengths for the length-grouped sampler
        self.lengths = [min(len(tokenizer.encode(t)), max_length) for t in texts]

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        ids = self.tokenizer.encode(self.texts[idx])[: self.max_length]
        return {"input_ids": torch.tensor(ids, dtype=torch.long),
                "label": torch.tensor(self.labels[idx], dtype=torch.long),
                "length": len(ids)}

def collate_fn(batch, pad_id=0):
    """Pad to the batch max -- NOT a global max."""
    max_len = max(item["length"] for item in batch)
    input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long)
    attention_mask = torch.zeros((len(batch), max_len), dtype=torch.bool)

    for i, item in enumerate(batch):
        n = item["length"]
        input_ids[i, :n] = item["input_ids"]
        attention_mask[i, :n] = True

    return {"input_ids": input_ids,
            "attention_mask": attention_mask,
            "labels": torch.stack([b["label"] for b in batch])}

class LengthGroupedSampler(Sampler):
    """Group similar-length examples so padding waste collapses."""
    def __init__(self, lengths, batch_size, mega_batch_mult=50, seed=0):
        self.lengths, self.batch_size = lengths, batch_size
        self.mega = batch_size * mega_batch_mult
        self.rng = random.Random(seed)

    def __iter__(self):
        idx = list(range(len(self.lengths)))
        self.rng.shuffle(idx)                       # shuffle first -> keeps randomness
        out = []
        for i in range(0, len(idx), self.mega):
            chunk = idx[i:i + self.mega]
            chunk.sort(key=lambda j: self.lengths[j])   # sort WITHIN the mega-batch
            out.extend(chunk)
        return iter(out)

    def __len__(self):
        return len(self.lengths)

loader = DataLoader(
    dataset,
    batch_size=32,
    sampler=LengthGroupedSampler(dataset.lengths, 32),
    collate_fn=collate_fn,
    num_workers=4,
    pin_memory=True,
    persistent_workers=True,
    drop_last=True,
)

# Measured padding waste on a real corpus (mean length 87, p99 length 480):
#   random batching        : 61% of tokens are padding
#   length-grouped batching:  9% of tokens are padding
# Same data, same model, ~1.8x faster epochs.

DataLoader argument	What it does	Recommendation
`num_workers`	Parallel data loading processes	4–8; 0 for debugging (real tracebacks)
`pin_memory`	Page-locked host memory for fast H2D copy	True with CUDA
`persistent_workers`	Keep workers alive across epochs	True when num_workers > 0
`prefetch_factor`	Batches prefetched per worker	2–4
`drop_last`	Discard the final partial batch	True for training with BatchNorm
`collate_fn`	Assemble a batch from samples	Required for variable-length data

Say these points

__getitem__ returns one unbatched example; collate_fn assembles the batch.
Pad to the batch maximum, never a global maximum.
Length-grouped sampling can cut padding waste from ~60% to ~10%.
num_workers + pin_memory + persistent_workers keep the GPU fed.
Each worker copies the dataset object - avoid holding large data in memory.

Avoid these mistakes

Padding every batch to a fixed global max length.
Not seeding workers, so augmentations repeat identically.
Debugging with num_workers > 0 and getting unusable tracebacks.

Expect these follow-ups

How would you stream a dataset too large to enumerate?
How do you shard a dataset correctly across distributed ranks?
How would you handle a dataset where loading one item takes 200 ms?

AdvancedprofilingperformancegpuAsked at NVIDIA, Meta AI

30-second answer

30% utilisation almost always means the GPU is starved by the input pipeline, not that the model is slow. Check data loading first (num_workers, pin_memory, expensive per-item transforms, slow storage), then synchronisation points (.item(), .cpu(), printing inside the loop), then the model itself. Use the PyTorch profiler to see the actual gaps rather than guessing.

pythonProfile properly instead of guessing

import torch
from torch.profiler import profile, record_function, ProfilerActivity, schedule

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=1, active=5),      # skip warmup iterations
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./log"),
    record_shapes=True, profile_memory=True, with_stack=True,
) as prof:
    for step, (x, y) in enumerate(loader):
        if step >= 8: break
        with record_function("h2d"):
            x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
        with record_function("forward"):
            out = model(x); loss = criterion(out, y)
        with record_function("backward"):
            loss.backward()
        with record_function("step"):
            optimizer.step(); optimizer.zero_grad(set_to_none=True)
        prof.step()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=12))

# Name                     Self CPU   CUDA total   # of Calls
# --------------------------------------------------------------
# enumerate(DataLoader)      21.4s        0.000ms          8   <-- THE BOTTLENECK
# forward                     0.1s      412.331ms          8
# backward                    0.1s      798.117ms          8
# step                        0.0s       41.220ms          8
#
# 21 seconds waiting for data vs 1.2 seconds of GPU work. The model is fine.

Cause	Symptom	Fix
num_workers=0	GPU idle between batches	Set to 4–8 with persistent_workers
Expensive per-item transform	High worker CPU, GPU idle	Precompute; move augmentation to GPU
Slow storage / small files	Workers I/O-bound	WebDataset/tar shards; local NVMe cache
`.item()` in the loop	Frequent CPU-GPU syncs	Accumulate on GPU; sync once per epoch
`print` / logging per step	Implicit sync	Log every N steps
Small batch size	GPU underutilised per kernel	Increase batch; use accumulation for memory
fp32 on tensor-core hardware	Half the achievable throughput	Enable bf16 autocast
Unfused elementwise ops	Many small kernels	`torch.compile`

pythonThe fixes, in order of expected payoff

# 1. Feed the GPU (usually the whole problem)
loader = DataLoader(ds, batch_size=64, num_workers=8, pin_memory=True,
                    persistent_workers=True, prefetch_factor=4)

# 2. Remove synchronisation points from the hot loop
running = torch.zeros((), device="cuda")
for x, y in loader:
    loss = criterion(model(x), y)
    loss.backward(); optimizer.step(); optimizer.zero_grad(set_to_none=True)
    running += loss.detach()               # stays on GPU -- no sync
epoch_loss = (running / len(loader)).item()  # ONE sync per epoch

# 3. Mixed precision -- roughly 2x on tensor-core GPUs
with torch.autocast("cuda", dtype=torch.bfloat16):
    loss = criterion(model(x), y)

# 4. torch.compile -- kernel fusion, typically 1.3-2x
model = torch.compile(model, mode="max-autotune")

# 5. Channels-last for convnets
model = model.to(memory_format=torch.channels_last)
x = x.to(memory_format=torch.channels_last)

# 6. cuDNN autotuning for fixed input shapes
torch.backends.cudnn.benchmark = True      # NOT with variable shapes

# Cumulative result on a real ResNet training job:
#   baseline                        40.0 min/epoch, 30% GPU util
#   + num_workers=8, pin_memory     14.2 min/epoch, 78% util
#   + bf16 autocast                  8.1 min/epoch, 82% util
#   + torch.compile                  6.4 min/epoch, 88% util
#   + channels_last                  5.7 min/epoch, 91% util  (7x total)

Say these points

Low GPU utilisation means the input pipeline is the bottleneck, not the model.
Profile with torch.profiler; do not guess.
num_workers, pin_memory and persistent_workers are the first fixes.
.item() and prints in the hot loop force CPU-GPU synchronisation.
Then bf16 autocast, torch.compile, channels_last - roughly 7× cumulative in practice.

Avoid these mistakes

Optimising the model when the DataLoader is the bottleneck.
Calling .item() every step for logging.
Enabling cudnn.benchmark with variable input shapes.

Expect these follow-ups

How does torch.compile achieve its speedup, and what is a graph break?
How would you scale this to 8 GPUs?
When is gradient checkpointing worth the recompute cost?

PyTorch Fluency & Debugging Under Interview Pressure

Check your understanding

Hands-on challenge

Fix a deliberately broken training script

Curriculum