AI infrastructure engineer coding interview questions

As asked

You are training a model with 7B parameters using DDP across 64 GPUs. Calculate the per-GPU data transmitted during one allreduce of the full gradient, assuming bf16 precision. How does this change if you switch to FSDP FULL_SHARD?

Sample answer outline

DDP allreduce: each GPU sends and receives 2 * (N-1)/N * message_size; for N=64 this is approximately 2 * message_size. 7B params * 2 bytes (bf16) = 14 GB per allreduce. With FSDP FULL_SHARD, the reduce-scatter transmits message_size * (N-1)/N and all-gather transmits the same, totaling 2 * 14 GB / 64 * (N-1) roughly but each rank only holds 1/N of the parameters, so peak in-flight per rank is much smaller. The candidate should derive the formula not just give the number.

Reference implementation (python)

# Estimate allreduce traffic for DDP vs FSDP
num_params = 7e9
bytes_per_param_bf16 = 2  # bf16
world_size = 64

# DDP: ring-allreduce sends 2*(N-1)/N * data per GPU
ddp_bytes_per_gpu = 2 * (world_size - 1) / world_size * num_params * bytes_per_param_bf16
print(f"DDP allreduce per GPU: {ddp_bytes_per_gpu / 1e9:.2f} GB")

# FSDP FULL_SHARD: reduce-scatter + all-gather, each ~data/N per rank per step
# Each rank stores 1/N params, but communicates full gradient volume in sharded fashion
fsdp_reduce_scatter = num_params * bytes_per_param_bf16 * (world_size - 1) / world_size
fsdp_all_gather = num_params * bytes_per_param_bf16 * (world_size - 1) / world_size
print(f"FSDP reduce-scatter per GPU: {fsdp_reduce_scatter / 1e9:.2f} GB")

Expect these follow-ups

• How would you measure actual allreduce bandwidth vs. theoretical peak bandwidth in your cluster?

As asked

Given a kernel that performs 512 GFLOPS of compute and reads/writes 8 GB of data, is it compute-bound or memory-bandwidth-bound on an A100 SXM? Show your calculation and explain what the roofline model predicts for the optimal kernel runtime.

Sample answer outline

A100 SXM peak bf16 TFLOPS is 312 TFLOPS; peak HBM2e bandwidth is 2 TB/s. Arithmetic intensity = 512 GFLOPS / 8 GB = 64 FLOPS/byte. Ridge point = 312e12 / 2e12 = 156 FLOPS/byte. Since 64 < 156, the kernel is memory-bandwidth-bound. Roofline-predicted time = 8 GB / 2 TB/s = 4 ms. A strong candidate should also discuss how to improve arithmetic intensity via kernel fusion.

Reference implementation (python)

# Roofline model calculation
peak_tflops = 312e12  # A100 SXM bf16
peak_bandwidth = 2e12  # bytes/s, A100 HBM3

kernel_flops = 512e9
kernel_bytes = 8e9

arithmetic_intensity = kernel_flops / kernel_bytes  # FLOPS/byte
ridge_point = peak_tflops / peak_bandwidth

if arithmetic_intensity < ridge_point:
    bound = "memory-bandwidth"
    roofline_time = kernel_bytes / peak_bandwidth
else:
    bound = "compute"
    roofline_time = kernel_flops / peak_tflops

print(f"Arithmetic intensity: {arithmetic_intensity:.1f} FLOPS/byte")
print(f"Ridge point: {ridge_point:.1f} FLOPS/byte")
print(f"Bound: {bound}")
print(f"Roofline predicted time: {roofline_time*1000:.2f} ms")

Expect these follow-ups

• If you fused two consecutive elementwise operations on this kernel, how would the arithmetic intensity change?

As asked

Write a Python script that samples nvidia-smi every 5 seconds and logs per-GPU utilization and memory usage to a CSV, with timestamps. The script should handle the case where a GPU disappears from the listing (e.g., after a driver reset).

Sample answer outline

The script should use subprocess to run nvidia-smi --query-gpu and --format=csv with a loop, parse the CSV output, handle missing or empty GPU rows gracefully, and write a timestamped CSV. A strong answer uses nvidia-smi's built-in --loop flag or subprocess.Popen for streaming, handles subprocess.CalledProcessError on driver failure, and notes that pynvml is a more reliable Python interface than parsing text output.

Reference implementation (python)

import subprocess, csv, time, datetime, sys

FIELDS = ["index", "utilization.gpu", "memory.used", "memory.total"]

def sample_gpus():
    try:
        result = subprocess.run(
            ["nvidia-smi", f"--query-gpu={','.join(FIELDS)}",
             "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=10
        )
        rows = []
        for line in result.stdout.strip().splitlines():
            vals = [v.strip() for v in line.split(",")]
            if len(vals) == len(FIELDS):
                rows.append(dict(zip(FIELDS, vals)))
        return rows
    except Exception as e:
        print(f"nvidia-smi error: {e}", file=sys.stderr)
        return []

with open("gpu_log.csv", "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=["timestamp"] + FIELDS)
    writer.writeheader()
    while True:
        ts = datetime.datetime.utcnow().isoformat()
        for row in sample_gpus():
            writer.writerow({"timestamp": ts, **row})
        f.flush()
        time.sleep(5)

Expect these follow-ups

• How would you replace this with pynvml to avoid the overhead of forking a subprocess every 5 seconds?

As asked

Write a function that takes a model config (num_layers, num_heads, head_dim, dtype_bytes) and serving config (max_seq_len, max_batch_size) and returns the peak KV cache memory in bytes. Then compute it for Llama-2-70B in fp16 serving 32 concurrent requests at 4096 tokens each.

Sample answer outline

KV cache = 2 * num_layers * num_heads * head_dim * dtype_bytes * max_seq_len * max_batch_size. For Llama-2-70B: 80 layers, 64 heads (GQA uses 8 KV heads), head_dim=128, fp16=2 bytes, 4096 tokens, 32 requests. With GQA the answer must use num_kv_heads (8) not num_heads. This is 2 * 80 * 8 * 128 * 2 * 4096 * 32 = approximately 42 GB. A strong candidate mentions GQA/MQA and its memory savings.

Reference implementation (python)

def kv_cache_bytes(
    num_layers: int,
    num_kv_heads: int,  # use KV heads for GQA/MQA
    head_dim: int,
    dtype_bytes: int,
    max_seq_len: int,
    max_batch_size: int,
) -> int:
    # Factor of 2 for K and V
    return 2 * num_layers * num_kv_heads * head_dim * dtype_bytes * max_seq_len * max_batch_size

# Llama-2-70B with GQA (8 KV heads), fp16, 4096 tokens, 32 concurrent
bytes = kv_cache_bytes(
    num_layers=80,
    num_kv_heads=8,   # GQA reduces from 64 to 8
    head_dim=128,
    dtype_bytes=2,    # fp16
    max_seq_len=4096,
    max_batch_size=32,
)
print(f"KV cache: {bytes / 1e9:.1f} GB")

Expect these follow-ups

• How does grouped-query attention (GQA) reduce KV cache size, and what is the tradeoff in model quality?

As asked

You have N=4 ranks. Ranks 0, 1, and 3 call dist.all_reduce(), but rank 2 mistakenly calls dist.broadcast() with the same tensor. Write a minimal reproducer and describe what you observe and how you fix it.

Sample answer outline

The collective type mismatch causes ranks to send incompatible protocol headers to each other; NCCL will either hang indefinitely or raise an internal error after a timeout. The reproducer should launch 4 processes with torch.multiprocessing or torchrun, have rank 2 call a different collective, and demonstrate the hang. Fix by ensuring all ranks call the same collective with the same arguments. Defensive pattern: use TORCH_DISTRIBUTED_DEBUG=DETAIL to log collective call sites.

Reference implementation (python)

import torch
import torch.distributed as dist
import os

def worker(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    tensor = torch.ones(4).cuda(rank)
    
    if rank == 2:
        # BUG: rank 2 calls broadcast instead of all_reduce
        dist.broadcast(tensor, src=0)  # hangs or errors
    else:
        dist.all_reduce(tensor)  # all others call all_reduce
    
    print(f"Rank {rank} done: {tensor}")

# Fix: all ranks must call the same collective
# dist.all_reduce(tensor) on every rank

if __name__ == "__main__":
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12355"
    torch.multiprocessing.spawn(worker, args=(4,), nprocs=4)

Expect these follow-ups

• How would you add an assertion in a multi-rank training loop to verify all ranks are about to call the same collective type?

As asked

When using gradient accumulation with PyTorch DDP and accumulating over K micro-batches before each optimizer step, what is the incorrect pattern that causes unnecessary allreduce calls on every micro-batch, and how do you fix it?

Sample answer outline

DDP triggers allreduce on every backward() call by default. Without no_sync(), every micro-batch backward triggers an allreduce, wasting communication for intermediate accumulation steps. Fix: wrap all micro-batch backward passes except the last in model.no_sync() context manager. The last backward triggers the real allreduce. Also need to scale gradients by 1/K or loss by 1/K before accumulation to match full-batch semantics.

Reference implementation (python)

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

# WRONG: triggers allreduce on every micro-batch
def train_wrong(model, batches, optimizer):
    for i, batch in enumerate(batches):
        loss = model(batch) / len(batches)
        loss.backward()  # allreduce fires here every time!
    optimizer.step()
    optimizer.zero_grad()

# CORRECT: suppress allreduce on intermediate micro-batches
def train_correct(model: DDP, batches, optimizer):
    for i, batch in enumerate(batches):
        is_last = (i == len(batches) - 1)
        ctx = torch.distributed.algorithms.ddp_comm_hooks.default_hooks
        # Use no_sync() to suppress gradient sync on intermediate steps
        if not is_last:
            with model.no_sync():
                loss = model(batch) / len(batches)
                loss.backward()
        else:
            loss = model(batch) / len(batches)
            loss.backward()  # allreduce fires here only
    optimizer.step()
    optimizer.zero_grad()

Expect these follow-ups

• Does using FSDP require a different pattern for gradient accumulation compared to DDP?

As asked

Write a simulation that models how increasing batch size affects per-token latency and throughput for an LLM decode step, assuming decode is memory-bandwidth-bound. Show how you would determine the optimal batch size given a latency SLA of 50 ms per token.

Sample answer outline

In the memory-bandwidth-bound regime, decode time per step scales linearly with batch size (each sequence reads the full model weights once). Throughput = batch_size / latency_per_step, which is flat in the bandwidth-bound regime since latency also scales linearly. The model should show that throughput saturates once compute-bound, and the optimal batch size for a 50 ms SLA is bandwidth / (model_bytes / 0.05). A strong answer distinguishes prefill (compute-bound) from decode (memory-bound).

Reference implementation (python)

def simulate_decode_throughput(
    model_bytes: float,         # total model parameter bytes
    hbm_bandwidth: float,       # bytes/sec (e.g. 2e12 for A100)
    max_batch: int = 64,
    latency_sla_ms: float = 50,
):
    results = []
    for bs in range(1, max_batch + 1):
        # Each decode step reads model weights once per batch
        # Time scales linearly with batch size in BW-bound regime
        latency_ms = (model_bytes / hbm_bandwidth) * 1000 * bs
        throughput_tok_per_sec = bs / (latency_ms / 1000)
        meets_sla = latency_ms <= latency_sla_ms
        results.append((bs, latency_ms, throughput_tok_per_sec, meets_sla))
    return results

# 7B model in fp16 = 14 GB, A100 HBM = 2 TB/s
for bs, lat, tput, sla in simulate_decode_throughput(14e9, 2e12, latency_sla_ms=50)[:10]:
    print(f"bs={bs:3d}  lat={lat:.1f}ms  tput={tput:.0f} tok/s  sla={'OK' if sla else 'MISS'}")

Expect these follow-ups

• How does tensor parallelism across multiple GPUs change the bandwidth-bound vs. compute-bound boundary?

As asked

Implement a simple ring-allreduce for a list of tensors across N ranks. Each rank should send its chunk to the next rank and receive from the previous rank, accumulating partial reductions. Show the reduce-scatter phase only.

Sample answer outline

The implementation should split the tensor into N chunks, iterate N-1 times sending the current chunk to rank+1 and receiving from rank-1, accumulating the received chunk into the local buffer. Key correctness points: non-blocking sends/receives to avoid deadlock, correct wraparound for rank 0 and rank N-1, and proper chunk indexing. A strong candidate notes that a real implementation uses CUDA-aware MPI or NCCL and avoids this Python approach in production.

Reference implementation (python)

import torch
import torch.distributed as dist

def ring_reduce_scatter(tensor: torch.Tensor, rank: int, world_size: int) -> torch.Tensor:
    """Reduce-scatter phase of ring-allreduce."""
    chunks = tensor.chunk(world_size)
    # Each rank owns one chunk after reduce-scatter
    buf = chunks[rank].clone()
    
    send_rank = (rank + 1) % world_size
    recv_rank = (rank - 1) % world_size
    
    for step in range(world_size - 1):
        send_chunk_idx = (rank - step) % world_size
        recv_chunk_idx = (rank - step - 1) % world_size
        
        recv_buf = torch.empty_like(chunks[recv_chunk_idx])
        send_req = dist.isend(chunks[send_chunk_idx], dst=send_rank)
        recv_req = dist.irecv(recv_buf, src=recv_rank)
        send_req.wait(); recv_req.wait()
        
        # Accumulate into the chunk we will own
        if recv_chunk_idx == rank:
            buf += recv_buf
    
    return buf

Expect these follow-ups

• How would you extend this to the all-gather phase to complete a full allreduce?

As asked

Write a function that takes pipeline depth P and number of micro-batches M and computes the bubble fraction (fraction of time the pipeline is not fully occupied) using the 1F1B (one forward one backward) schedule.

Sample answer outline

Bubble fraction = (P - 1) / M for the 1F1B schedule. The function should verify M >= P for a valid schedule, compute the formula, and optionally visualize the schedule as a grid. A strong candidate knows this comes from the Megatron paper and explains why increasing M reduces the bubble at the cost of more micro-batch overhead.

Reference implementation (python)

def pipeline_bubble_fraction(pipeline_depth: int, num_microbatches: int) -> float:
    """
    Compute the pipeline bubble fraction for the 1F1B schedule.
    
    Args:
        pipeline_depth: number of pipeline stages P
        num_microbatches: number of micro-batches M (must be >= P)
    
    Returns:
        bubble_fraction in [0, 1)
    """
    P, M = pipeline_depth, num_microbatches
    if M < P:
        raise ValueError(f"M ({M}) must be >= P ({P}) for 1F1B to be valid")
    # Warm-up phase has (P-1) idle slots; total slots = M * P
    # Bubble slots = P - 1 (before pipeline is full) + P - 1 (drain)
    bubble_fraction = (P - 1) / M
    return bubble_fraction

# Example: 8 stages, 16 micro-batches
print(pipeline_bubble_fraction(8, 16))  # 0.4375
print(pipeline_bubble_fraction(8, 64))  # 0.109375

Expect these follow-ups

• How does the interleaved 1F1B schedule (virtual pipeline stages) improve the bubble fraction?

As asked

Implement a simplified forward pass of LayerNorm in OpenAI Triton that fuses the mean, variance, and normalization into a single kernel. The input is a 2D tensor of shape (N, D) and normalization is over the D dimension.

Sample answer outline

The Triton kernel should use tl.program_id to identify the row, load D elements in blocks, compute mean via tl.sum/D, compute variance similarly, then normalize in one pass, writing output. Key points: use tl.constexpr for BLOCK_D, handle non-power-of-two D with masking, and accumulate in fp32 for numerical stability before casting back to the input dtype. A strong answer discusses when this is faster than torch.nn.LayerNorm (on small D where the HBM round trip costs dominate).

Reference implementation (python)

import triton
import triton.language as tl
import torch

@triton.jit
def layernorm_fwd_kernel(
    X, Y, Mean, Rstd,
    stride, N, eps,
    BLOCK_D: tl.constexpr,
):
    row = tl.program_id(0)
    X_ptr = X + row * stride
    Y_ptr = Y + row * stride
    cols = tl.arange(0, BLOCK_D)
    mask = cols < N
    x = tl.load(X_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    mean = tl.sum(x, axis=0) / N
    var = tl.sum((x - mean) * (x - mean), axis=0) / N
    rstd = 1.0 / tl.sqrt(var + eps)
    y = (x - mean) * rstd
    tl.store(Y_ptr + cols, y.to(tl.float16), mask=mask)
    tl.store(Mean + row, mean)
    tl.store(Rstd + row, rstd)

Expect these follow-ups

• How would you add the weight and bias (gamma, beta) parameters to this kernel without adding extra HBM reads?