ai_ml / ml frameworks / 04_pytorch_interview.md

PyTorch — Common Interview Questions and Answers

4 interview angles 9 min read source

PyTorch — Common Interview Questions and Answers

PyTorch is an open-source deep learning framework from Meta AI with a Pythonic, define-by-run programming model. Popular because:

  • Dynamic computation graphs — Python control flow works naturally.
  • Debuggability — standard Python stack traces, breakpoints anywhere.
  • NumPy-like tensor API — small learning curve.
  • Hugging Face / LLM ecosystem — almost all modern models are PyTorch-first.

2. What’s a tensor and how does it differ from a NumPy array?

A torch.Tensor is an n-dimensional array that adds:

  • GPU acceleration via .to("cuda") / .cuda().
  • Autograd via requires_grad=True — tracks operations for gradient computation.
  • Deferred CUDA execution — GPU ops are async; synchronize with torch.cuda.synchronize() for benchmarking.

You can zero-copy convert: torch.from_numpy(arr) and tensor.numpy().


3. How does autograd work?

PyTorch tracks operations on tensors with requires_grad=True into a dynamic computation graph (built fresh each forward pass). When you call .backward() on a scalar output, PyTorch walks the graph backward, applying the chain rule, and populates .grad on leaf tensors.

x = torch.tensor(2.0, requires_grad=True)
y = x ** 3 + 2 * x
y.backward()
print(x.grad)   # 14

Gradients accumulate by default — that’s why you call optimizer.zero_grad() at each step.


4. What is nn.Module?

Base class for all neural network components. Two key behaviors:

  • Parameter registration — assigning an nn.Parameter or another nn.Module as an attribute automatically registers it so model.parameters() enumerates everything.
  • Forward definition — you override forward(self, x); calling model(x) invokes forward with hooks.
class Net(nn.Module):
    def __init__(self):
        super().__init__()                      # required
        self.fc = nn.Linear(10, 1)              # registered automatically

    def forward(self, x):
        return self.fc(x)

5. What does optimizer.zero_grad() do?

Clears .grad on all parameters tracked by the optimizer. Necessary because PyTorch accumulates gradients across .backward() calls — without zeroing, each step would train on the sum of all prior gradients.

The accumulation behavior is intentional: it enables gradient accumulation (train as if batch size were N×, on hardware that only fits batch N).


6. What’s the difference between model.train() and model.eval()?

Toggles modules that behave differently at train vs eval time:

  • Dropout — active in train mode (randomly zeros neurons); identity in eval.
  • BatchNorm — uses batch stats in train mode; running stats (frozen) in eval.

Always call model.eval() before inference and validation. Pair with torch.no_grad() to also disable autograd tracking:

model.eval()
with torch.no_grad():
    preds = model(x_val)

7. What’s torch.no_grad() and when do you use it?

Context manager that disables autograd tracking. Saves memory (no graph built) and slightly speeds up forward passes. Use for inference, validation, and any code where you don’t need gradients.

with torch.no_grad():
    logits = model(x)

Alternative: @torch.inference_mode() (slightly faster, more restrictive — tensors created inside can’t be used in autograd later).


8. How do you move a model and data to GPU?

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)  # AFTER .to()

for xb, yb in loader:
    xb, yb = xb.to(device), yb.to(device)
    ...

Gotchas:

  • Create the optimizer after moving the model to GPU.
  • Both data and model must be on the same device for any op (mismatch raises RuntimeError: Expected all tensors to be on the same device).
  • .cuda() is shorthand for .to("cuda").

9. What’s a Dataset and DataLoader?

  • Dataset — defines how to access individual samples. Subclass implements __len__ and __getitem__.
  • DataLoader — wraps a Dataset and handles batching, shuffling, multi-process loading, prefetching.
class MyDataset(Dataset):
    def __init__(self, data): self.data = data
    def __len__(self): return len(self.data)
    def __getitem__(self, i): return self.data[i]

loader = DataLoader(MyDataset(data), batch_size=64, shuffle=True, num_workers=4)

num_workers > 0 enables multiprocessing data loading. pin_memory=True speeds GPU transfer. On Windows, requires if __name__ == "__main__": guard.


10. How do you save and load a model?

Save the state_dict (parameter tensors only), not the whole model:

torch.save(model.state_dict(), "model.pt")

# Loading
model = Net()                          # must instantiate first
model.load_state_dict(torch.load("model.pt", map_location="cpu"))
model.eval()

Why state_dict and not torch.save(model)? The latter pickles the whole class — fragile across code refactors. state_dict is just the tensors and survives class changes.

For full training resume (optimizer state, epoch, etc.):

torch.save({
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}, "checkpoint.pt")

11. What’s the difference between view, reshape, and permute?

Op Memory Notes
view(shape) requires contiguous tensor fails if not contiguous; cheapest
reshape(shape) may copy if not contiguous safer; copies if needed
permute(*dims) does not copy swaps dimensions; tensor becomes non-contiguous
transpose(d1, d2) does not copy swaps two dims; non-contiguous after
contiguous() copies if needed call before view if tensor is non-contiguous
x.permute(0, 2, 1).contiguous().view(B, -1)

Common bug: view after transposeRuntimeError: view size is not compatible. Fix with .contiguous() or use reshape.


12. What’s broadcasting?

Implicit expansion of mismatched tensor shapes when their dimensions are compatible. Rules (right-aligned):

  • Dimensions are equal, or
  • One of them is 1, or
  • One doesn’t exist.
a = torch.zeros(3, 1, 5)
b = torch.zeros(   4, 1)   # broadcasts to (1, 4, 1)
c = a + b                  # shape (3, 4, 5)

Same rules as NumPy. Source of subtle bugs when shapes accidentally broadcast in ways you didn’t intend.


13. How do you implement a custom loss function?

Two ways:

# As a function
def my_loss(pred, target):
    return ((pred - target) ** 2).mean()

# As a module (better if it has state / parameters)
class MyLoss(nn.Module):
    def __init__(self, weight=1.0):
        super().__init__()
        self.weight = weight

    def forward(self, pred, target):
        return self.weight * ((pred - target) ** 2).mean()

Anything you write with tensor ops gets autograd for free. Only override torch.autograd.Function if you need a custom backward (e.g., straight-through estimator, custom CUDA op).


14. What’s mixed-precision training?

Train using float16 or bfloat16 for most ops (memory + speed wins) and float32 for numerically sensitive parts (loss scaling, optimizer state). PyTorch has torch.cuda.amp:

scaler = torch.cuda.amp.GradScaler()
for xb, yb in loader:
    optimizer.zero_grad()
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        logits = model(xb)
        loss = loss_fn(logits, yb)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

On modern hardware (A100, H100), use bfloat16 — wider dynamic range, no loss scaling needed. On older GPUs, float16 with GradScaler.


15. How do you do distributed training?

The two main options:

  • DistributedDataParallel (DDP) — data parallelism. Each GPU has a full model copy; gradients are all-reduced. Standard for fitting model-on-one-GPU at higher batch sizes.
  • FullyShardedDataParallel (FSDP) — shards model parameters, gradients, optimizer state across GPUs. For models that don’t fit on one GPU. Like DeepSpeed ZeRO-3.
torchrun --nproc_per_node=8 train.py
# inside train.py:
torch.distributed.init_process_group("nccl")
model = DDP(model, device_ids=[local_rank])

Higher-level abstractions: Hugging Face accelerate (wraps DDP/FSDP/DeepSpeed in one API), PyTorch Lightning, DeepSpeed.


16. What’s torch.compile?

PyTorch 2.0 feature that JIT-compiles model code into optimized fused kernels (Inductor backend, often via Triton on GPU). 2-3x speedups common for many models without changing code:

model = torch.compile(model)            # that's it

First forward pass is slow (compilation); subsequent calls run the compiled graph. Falls back to eager on unsupported ops (“graph breaks”). The TF @tf.function analog.


17. Common training bugs?

  • Forgot optimizer.zero_grad() — gradients accumulate; loss explodes or learning fails.
  • Wrong loss reductionCrossEntropyLoss returns mean by default; if you sum manually, scale accordingly.
  • Targets wrong dtypeCrossEntropyLoss wants long (int64) class indices, not one-hot floats.
  • view after transpose — non-contiguous tensor; use reshape or .contiguous().
  • Inputs not on GPURuntimeError: Expected all tensors to be on the same device.
  • model.eval() forgotten — dropout/batchnorm misbehave at inference.
  • shuffle=True on validation loader — makes per-epoch metrics non-comparable.
  • DataLoader worker hang on Windows — missing if __name__ == "__main__":.
  • Memory leak from holding references to tensor graphs — use .detach() or .item() when storing loss values for logging.

18. What’s .detach() vs .item()?

  • .detach() — returns a new tensor sharing storage, but disconnected from the autograd graph. Used to stop gradient flow without copying data.
  • .item() — converts a 0-dim tensor to a Python scalar. Use for logging scalar loss values (avoids holding the autograd graph alive).
running_loss += loss.item()             # right — Python float
running_loss += loss                    # wrong — keeps graph alive across iterations, OOM

19. How do you debug NaN losses?

  • Check learning rate — too high often causes divergence.
  • Check inputs — NaN/Inf in data leaks through.
  • Use torch.autograd.set_detect_anomaly(True) — slow but pinpoints which op produced NaN.
  • Add gradient clippingtorch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0).
  • Stability trickslog_softmax + nll_loss instead of log(softmax); use from_logits paths; avoid log(0) via log(x + eps).
  • Mixed precision — fp16 can underflow; switch to bf16 or use GradScaler.

20. How do you deploy a PyTorch model?

Several paths depending on target:

Target Path
Python server TorchServe, FastAPI + model loaded in memory
LLM serving vLLM, TensorRT-LLM, Triton Inference Server
Cross-runtime export to ONNX, run with ONNX Runtime / TensorRT
Mobile ExecutorTorch (newer), PyTorch Mobile (legacy)
Browser export to ONNX → onnx.js, or compile via WebAssembly
Quantized torch.ao.quantization (int8/int4), then deploy

For production LLM serving, vLLM is the de facto standard (continuous batching, PagedAttention, OpenAI-compatible API). For traditional models, TorchServe or just FastAPI wrapping model(x).


Quick reference: the mental model

Concept Quick definition
Tensor n-d array, GPU-capable, optionally autograd-tracked
nn.Module building block; auto-registers params
nn.Parameter tensor marked as a trainable parameter
Autograd dynamic graph; .backward() populates .grad
optimizer.step() applies .grad to update parameters
zero_grad() clear .grad (else accumulates)
model.train()/.eval() toggle dropout/batchnorm mode
no_grad() disable autograd tracking (inference/eval)
DataLoader batch+shuffle+prefetch over Dataset
torch.compile JIT-compile model to optimized kernels
DDP / FSDP data / sharded distributed training

Interview angle

  • “Why did PyTorch win?” - define-by-run. The graph is built as Python executes, so debugging is ordinary Python with real stack traces and breakpoints. Combined with the research ecosystem, that made it the default and the target of essentially every published model.
  • “What does torch.no_grad() do and when do you need it?” - disables autograd graph construction, cutting memory and time. Use it for all inference and evaluation; forgetting it builds a graph you never backpropagate through and can exhaust memory.
  • “Why optimizer.zero_grad()?” - gradients accumulate across backward passes by default. Omitting it means each step uses the running sum of all previous gradients. The accumulation is deliberate, so you can simulate large batches across micro-batches.
  • “What is torch.compile?” - graph capture plus kernel fusion for speedups without changing model code. A first-call compilation cost, and it can fall back to eager on unsupported constructs, so measure rather than assume.