TensorFlow/PyTorch Framework Fundamentals Questions

Practical knowledge of a major deep learning framework. Includes understanding tensors, operations, building neural network layers, constructing models, and training loops. Ability to read and modify existing code in these frameworks. Knowledge of how to work with pre-built layers and models.

HardSystem Design

45 practiced

Design an inference optimization pipeline to deploy a large Transformer for low-latency GPU inference. Outline steps including exporting to ONNX or TorchScript, operator fusion, TensorRT conversion, dynamic batching, quantization options, server choices (TorchServe/TensorFlow Serving/TensorRT Inference Server), and how to measure/validate latency and correctness.

Sample Answer

Requirements & constraints:- Target: large Transformer (e.g., 7B–70B) on GPU for low-latency (p50/p90 targets), memory-limited, support variable-length inputs and attention KV cache, throughput vs latency trade-off.

High-level pipeline (steps):1. Clarify model I/O and runtime behaviors - Determine dynamic axes (batch, seq_len), autoregressive vs encoder-only, use of KV cache, beam search or sampling. - Define latency targets (p50/p90), acceptable accuracy loss (e.g., <=0.5% dist from FP32), max GPU memory.

2. Export - Prefer TorchScript for PyTorch-specific ops if staying in PyTorch runtime; otherwise export to ONNX (opset >= 13) for tool interoperability. - Ensure correct tracing vs scripting: script modules with control-flow; trace only pure feedforward graph parts. Mark KV cache as explicit inputs/outputs.

3. Graph-level optimizations / operator fusion - Run framework optimizers (torch.jit.freeze, torch.jit.optimized) then ONNX simplifier (onnx-simplifier). - Fuse common patterns: layernorm -> linear, add+gelu, attention matmuls; use vendor passes (ONNX Runtime graph optimizations, NVIDIA's torch-TRT or TensorRT plugins). - Replace inefficient ops (e.g., Softmax+Mask) with fused plugins supporting masked attention.

4. Convert to TensorRT (preferred for NVIDIA GPUs) - Use ONNX -> TensorRT conversion or torch-TRT/Torch-TensorRT for direct conversion. - Enable FP16 first (minimal accuracy loss), evaluate INT8 with calibration dataset for aggressive quantization. For INT8 use per-channel quantization and calibration (representative dataset). - For attention/KV cache, ensure engine supports dynamic shapes and explicit batch/seq dims (use optimization profiles).

5. Dynamic batching & shape management - Use dynamic shape support and multiple optimization profiles in TensorRT for different seq_len/batch combos. - Implement server-side dynamic batching (aggregate small requests within latency budget) and client-side adaptive batching based on arrival rates. - For autoregressive generation, use token-level batching (micro-batching) and batched KV cache updates.

6. Serving choices - NVIDIA Triton Inference Server (recommended): native TensorRT backends, dynamic batching, concurrent model instances, metrics (Prometheus), sequence models, and gRPC/HTTP. - TorchServe: simpler for TorchScript but less mature for TensorRT and multi-profile dynamic shapes. - TensorFlow Serving: use if model converted to SavedModel; less integrated for TensorRT optimizations. - Choose Triton when needing best GPU latency/throughput and operational telemetry.

7. Quantization options & trade-offs - FP16: best first step — often 1–2% negligible quality change, large speed/memory gains. - INT8: best throughput but risk of numeric degradation — do calibration and per-layer sensitivity analysis. Keep critical layers (softmax, layernorm) in FP16 if needed. - Mixed precision: keep attention softmax, layernorm in FP32/FP16 as required.

8. Validation: correctness & latency - Functional correctness: - Unit tests: compare outputs of reference FP32 model vs optimized model on a representative dataset. Use elementwise tolerance and statistical checks (mean/STD of diffs). - Behavioral tests: end-to-end sampling outputs, n-gram overlap, BLEU/ROUGE or human evaluation for generation tasks. - Latency & performance: - Measure p50/p90/p99 latency and throughput under realistic load using load testing (wrk, locust, k6, custom client). - Profile with NVIDIA nsight, Triton metrics (GPU util, SM efficiency), nvprof/tegrastats. - Test cold-start, warm-up (layer JIT/fusion warm-up), and steady-state. Include tail latencies with multi-tenant noise. - Regression gating: set automated acceptance criteria (e.g., p90 < target and accuracy drop < threshold).

9. Production considerations - Autoscaling policy based on p90 latency and GPU utilization; use model-ensemble or multiple TensorRT engines for different seq-len profiles. - Cache frequent prompts/responses, use early-exit or smaller models for short/simple requests. - Observability: expose Prometheus metrics (latency histograms, GPU mem), logging, request tracing, and A/B testing for quality. - Safety & rollback: shadow testing and canary rollout before full traffic.

Trade-offs- TorchScript keeps Pythonless runtime but less cross-platform; ONNX+TensorRT gives maximum performance on NVIDIA GPUs.- Aggressive INT8 yields more speed but higher validation cost and potential quality loss.- Dynamic batching improves throughput but can increase tail latency unless bounded.

Summary checklist- Export deterministically (script vs trace) with KV cache explicit.- Run graph fusion then TensorRT conversion with optimization profiles.- Apply FP16 then INT8 with calibration and per-layer sensitivity.- Serve with Triton for dynamic batching and telemetry.- Validate via unit comparisons + behavioral tests and measure p50/p90/p99, throughput, and GPU profiling; gate deployments on accuracy and latency metrics.

MediumTechnical

38 practiced

You're fine-tuning a pre-trained transformer (Hugging Face) for text classification in PyTorch. Describe the steps to replace the classification head, freeze the base model initially, specify different learning rates for the head and the base, and then unfreeze and fine-tune with layer-wise learning rate decay. Include code snippets or param group examples.

Sample Answer

Brief approach: load a pretrained AutoModel, replace its classifier with a task-specific head, freeze the base (encoder) weights, train only the head with a higher LR, then unfreeze and apply layer-wise learning rate decay (LLRD) to fine-tune lower layers with smaller LRs.

Code example (PyTorch + Hugging Face):

python

from transformers import AutoModel, AutoConfig, AutoTokenizer
import torch.nn as nn

model_name = "bert-base-uncased"
config = AutoConfig.from_pretrained(model_name, num_labels=3)
base = AutoModel.from_pretrained(model_name, config=config)

# Replace classification head
class Classifier(nn.Module):
    def __init__(self, base, hidden_dim, num_labels):
        super().__init__()
        self.base = base
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(hidden_dim, num_labels)
    def forward(self, **kwargs):
        outputs = self.base(**kwargs)
        pooled = outputs.pooler_output if hasattr(outputs, "pooler_output") else outputs.last_hidden_state[:,0]
        return self.classifier(self.dropout(pooled))

model = Classifier(base, base.config.hidden_size, config.num_labels)

Freeze base initially:

python

for param in model.base.parameters():
    param.requires_grad = False

Optimizer with different LRs (head vs base):

python

import torch.optim as optim

head_params = [p for n,p in model.named_parameters() if p.requires_grad]
optimizer = optim.AdamW([
    {"params": head_params, "lr": 5e-5, "weight_decay": 0.01}
])

Unfreeze and apply LLRD:

python

# Unfreeze all
for param in model.base.parameters():
    param.requires_grad = True

# Build layer-wise groups: assume encoder.layers.* naming
def get_layer_id(name):
    # example for BERT: encoder.layer.{i}
    import re
    m = re.search(r"encoder\.layer\.(\d+)", name)
    return int(m.group(1)) if m else -1

base_lr = 2e-5
head_lr = 5e-5
decay = 0.8
param_groups = {}
for name, param in model.named_parameters():
    if not param.requires_grad:
        continue
    if "classifier" in name or "pooler" in name:
        key = "head"
    else:
        lid = get_layer_id(name)
        key = f"layer_{lid}" if lid>=0 else "embeddings"
    param_groups.setdefault(key, {"params": [], "lr": None})
    param_groups[key]["params"].append(param)

# assign lrs: deeper layers get smaller lr via decay^layer_index
for key in list(param_groups.keys()):
    if key == "head":
        param_groups[key]["lr"] = head_lr
    elif key == "embeddings":
        param_groups[key]["lr"] = base_lr * (decay ** (base.config.num_hidden_layers + 1))
    else:
        idx = int(key.split("_")[1])
        # layers numbered 0..L-1; apply decay based on distance from top
        depth = base.config.num_hidden_layers - 1 - idx
        param_groups[key]["lr"] = base_lr * (decay ** depth)

optimizer = optim.AdamW(list(param_groups.values()), lr=base_lr, weight_decay=0.01)

Notes:- Tune base_lr, head_lr, and decay (common: decay 0.8–0.95).- Use weight decay only on weights, not biases/LayerNorm.- For large models use gradient accumulation, mixed precision (AMP), and lower base_lr.- Validate layer name patterns for your model (RoBERTa/DeBERTa differ).

HardTechnical

44 practiced

Implement a custom PyTorch autograd.Function for an operation where the forward() returns y = x * sigmoid(x) and backward() returns the correct gradient analytically. Provide the class skeleton with forward and backward static methods, show how to call it in a model, and explain how to test it with torch.autograd.gradcheck.

Sample Answer

To implement a custom PyTorch autograd Function for y = x * sigmoid(x) (sometimes called Swish), you need to write a class with static forward and backward methods. The backward must return the analytical gradient: if s = sigmoid(x), dy/dx = s + x * s * (1 - s).

Implementation, usage in a model, and gradcheck:

python

import torch
from torch.autograd import Function

class SwishFunction(Function):
    @staticmethod
    def forward(ctx, x):
        # x: Tensor
        s = torch.sigmoid(x)
        ctx.save_for_backward(x, s)   # save tensors needed for backward
        return x * s

    @staticmethod
    def backward(ctx, grad_output):
        x, s = ctx.saved_tensors
        # derivative: s + x * s * (1 - s)
        grad_x = grad_output * (s + x * s * (1 - s))
        return grad_x  # must return one tensor per forward arg

Using the custom op inside a model:

python

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.lin = nn.Linear(in_features, out_features)

    def forward(self, x):
        z = self.lin(x)
        return SwishFunction.apply(z)  # call the custom Function

Testing gradients with torch.autograd.gradcheck:- gradcheck compares your backward to finite differences; inputs must be double precision and require_grad=True.- Use a small tensor, not too large.

python

from torch.autograd import gradcheck

# prepare input: double, requires_grad
x = torch.randn(3, 4, dtype=torch.double, requires_grad=True)
# wrap in tuple for gradcheck
test = gradcheck(SwishFunction.apply, (x,), eps=1e-6, atol=1e-5, rtol=1e-3)
print("gradcheck passed:", test)

Notes and reasoning:- Save both x and sigmoid(x) to compute gradient efficiently and avoid recomputing expensive ops.- gradcheck requires double precision because finite-difference approximations are sensitive to precision.- If forward had additional non-tensor args, backward should return None for them.- Edge cases: ensure numeric stability for large magnitude x (sigmoid handles this), and test with a few random seeds and shapes.

EasyTechnical

35 practiced

Explain how automatic differentiation works in PyTorch (autograd). In your answer, include: what requires_grad means, how the dynamic computation graph is built and used during backward(), and a short example where you compute gradients of y = 3*x^2 + 2*x for x being a tensor with requires_grad=True.

Sample Answer

Automatic differentiation in PyTorch (autograd) is an eager, reverse-mode system that records operations on tensors to build a dynamic computation graph during the forward pass and then uses that graph to compute gradients in the backward pass.

- requires_grad: a tensor attribute; if True, PyTorch will track all operations on that tensor and store the information needed to compute gradients w.r.t. it. Leaf tensors with requires_grad=True will accumulate gradients in their .grad field after backward().

- Dynamic computation graph: as you perform operations, autograd creates nodes (Operations) and edges (tensors) linking inputs to outputs. This graph is dynamic (rebuilt each forward pass), which makes PyTorch flexible for control flow. When you call backward() on a scalar output, autograd traverses the graph in reverse, applies the chain rule at each node, and accumulates gradients into .grad of leaf tensors.

- Notes: Use torch.no_grad() or .detach() to stop tracking. If the output is non-scalar, pass a gradient argument to backward().

Example:

python

import torch

# x is a leaf tensor that we want gradients for
x = torch.tensor(4.0, requires_grad=True)

# forward: y = 3*x^2 + 2*x
y = 3 * x**2 + 2 * x

# backward: compute dy/dx. y is scalar so we can call backward() directly
y.backward()

print(x.grad)  # expected derivative: dy/dx = 6*x + 2 -> 6*4 + 2 = 26.0

Key takeaways:- Autograd builds the graph during forward execution and uses reverse-mode AD to compute gradients.- requires_grad enables tracking; .grad holds accumulated gradients for leaf tensors after backward().- The graph is dynamic per forward pass, enabling flexible model logic.

EasyTechnical

46 practiced

In TensorFlow 2.x (eager mode), explain the difference between tf.constant and tf.Variable. Provide a minimal code example that creates a Variable x initialized to 0.0 and uses tf.GradientTape to perform one gradient descent step to minimize f(x) = (x - 5)^2. Explain how to update the variable and how the gradient is observed.

Sample Answer

tf.constant creates an immutable tensor whose value cannot change; tf.Variable creates a mutable tensor whose value can be updated (e.g., during training). In TF2 eager mode, Variables track gradients and are the objects you typically optimize.

Example: one GD step to minimize f(x) = (x-5)^2

python

import tensorflow as tf

# create a trainable variable
x = tf.Variable(0.0)  # initial value 0.0

lr = 0.1  # learning rate

with tf.GradientTape() as tape:
    y = (x - 5.0) ** 2  # compute loss

# compute gradient dy/dx
grad = tape.gradient(y, x)
print("gradient:", grad.numpy())  # should be -10.0 when x=0

# update variable: x = x - lr * grad
x.assign_sub(lr * grad)
print("updated x:", x.numpy())

Explanation:- GradientTape records operations on watched variables (Variables are watched by default), so tape.gradient returns d y / d x.- We update the Variable using assign_sub (in-place subtract) or assign(x - lr*grad). tf.constant cannot be updated, so gradients w.r.t. constants are not useful for optimization.- Use optimizers (e.g., tf.keras.optimizers.SGD) in real code for convenience and extra features (momentum, etc.).

Unlock Full Question Bank

Get access to hundreds of TensorFlow/PyTorch Framework Fundamentals interview questions and detailed answers.

Join thousands of developers preparing for their dream job.