Diagnose a PyTorch Training Bottleneck from Profiler Traces
Company: Imc
Role: Hardware Engineer
Category: Machine Learning
Difficulty: easy
Interview Round: Onsite
# Diagnose a PyTorch Training Bottleneck from Profiler Traces
You are given a PyTorch training program and profiler traces from representative steps. Identify the dominant bottleneck, explain which code path causes it, and propose a bounded change whose effect can be measured without changing model semantics unintentionally.
The trace may include CPU operators, CUDA kernels, memory copies, synchronization, data-loader work, and gaps between iterations. Explain how you distinguish input, host, device, communication, and memory bottlenecks.
### Constraints & Assumptions
- The code, profiler configuration, hardware, and trace are supplied during the exercise; do not assume a particular bottleneck.
- Warm-up, compilation, and initial data-loading steps may differ from steady state.
- Model quality and numerical behavior must be checked after a performance change.
- A faster isolated kernel is not sufficient if end-to-end step time does not improve.
### Clarifying Questions to Ask
- Is the target single-device, multi-device, or multi-node training?
- Which steps are warm-up, and is data loading included in the captured trace?
- Are mixed precision, compilation, gradient accumulation, and distributed data parallelism enabled?
- Which metric matters: examples per second, time to train, cost, device utilization, or tail step time?
- May batch size, model math, data order, or numerical precision change?
### What a Strong Answer Covers
- Correct trace scoping with steady-state steps and CPU-CUDA time alignment.
- Evidence for data starvation, host launch overhead, explicit or implicit synchronization, kernel inefficiency, memory pressure, or distributed communication.
- Attribution from a long span or idle gap back to the exact Python or operator path.
- One change at a time, with expected trace signature and an end-to-end before-and-after measure.
- Data-loader tuning, transfer overlap, batching, operator fusion, mixed precision, compilation, or communication overlap only when the trace supports it.
- Correctness checks for loss, gradients, convergence, sample order, precision, and memory behavior.
- Avoidance of profiler overhead and measurement artifacts.
### Follow-up Questions
1. GPU utilization is low and there are long gaps before each step. How do you separate data starvation from CPU launch overhead?
2. Why can calling `.item()` inside the training loop create a visible synchronization stall?
3. CUDA kernels are busy, but examples per second do not improve after a faster kernel. What else should you inspect?
4. How do you determine whether all-reduce is exposed or overlapped with backward computation?
5. Which correctness checks matter after enabling mixed precision or changing batch size?
Overview: Use PyTorch profiler traces to distinguish input starvation, CPU launch overhead, synchronization, CUDA kernels, memory limits, and distributed communication. The solution ties each optimization to a causal code path, predicted trace change, end-to-end throughput, and model-correctness checks.
Read the full Imc Hardware Engineer interview experience this question came from