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.

MediumTechnical

39 practiced

Describe how to convert a PyTorch model to TorchScript using torch.jit.trace and torch.jit.script. Provide code examples for both approaches, explain differences, and give examples of failures (e.g., data-dependent control flow) where scripting is required instead of tracing.

Sample Answer

Approach (brief): TorchScript converts PyTorch models to a statically analyzable format. Use torch.jit.trace when model control flow does not depend on input data (trace records tensor operations for a given example). Use torch.jit.script when model contains data-dependent Python control flow (if/for that depend on tensor values) or native Python logic — scripting compiles the code including control flow.

Tracing example:

python

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self): 
        super().__init__()
        self.lin = nn.Linear(10, 1)
    def forward(self, x):
        return torch.relu(self.lin(x))

m = MyModel().eval()
example = torch.randn(4, 10)
traced = torch.jit.trace(m, example)        # trace with example input
traced.save("m_traced.pt")
# load: traced = torch.jit.load("m_traced.pt")

Scripting example:

python

import torch
import torch.nn as nn

class ControlModel(nn.Module):
    def __init__(self): 
        super().__init__()
        self.lin = nn.Linear(10, 1)
    def forward(self, x):
        # control depends on data: different branch if batch mean > 0
        if x.mean() > 0:
            return self.lin(x) * 2
        else:
            return self.lin(x)
scripted = torch.jit.script(ControlModel())  # compiles control flow
scripted.save("m_scripted.pt")

Key differences and reasoning:- Tracing records the executed operator graph for the concrete example. It's simple and often faster to produce, but it will “lock in” the exact execution path seen during tracing; data-dependent branches not taken during the example are not represented.- Scripting parses and compiles your Python/torch code into TorchScript IR, preserving dynamic control flow that uses tensor values or Python logic. Use scripting when model behavior depends on inputs.

Examples of failures where scripting is required:- Data-dependent if/for that choose branches by tensor values (as in ControlModel). Tracing will capture only the branch exercised by the example input — other branches missing lead to incorrect behavior at runtime.- Loops whose iteration count depends on input tensors.- Usage of Python-side logic operating on tensor contents (e.g., converting tensor to Python bool via .item() or using tensor.sum() in an if).- Models using isinstance or dynamic attribute creation — scripting can often handle these; tracing cannot.

Practical tips:- Always test traced/scripted outputs against eager model on representative inputs (use torch.testing.assert_allclose).- For tracing, provide representative example inputs that exercise all relevant execution paths if possible (but that’s often impossible for data-dependent branching).- Prefer scripting for production models with any input-dependent control flow or Python logic; use tracing for simple, feed-forward nets for speed/compatibility.

HardTechnical

48 practiced

Outline how to implement pipeline parallelism for a very large model using torch.distributed.rpc or torch.distributed.pipeline.sync. Describe how to partition the model across worker ranks, schedule forward and backward passes, manage micro-batches to keep the pipeline filled, handle parameter synchronization and checkpointing, and measure the latency vs throughput trade-offs.

Sample Answer

Approach: Partition the huge model into sequential stages placed on different worker ranks, stream micro-batches through the stages (1F1B schedule) to overlap compute and keep GPUs busy, and combine pipeline parallelism with data-parallel parameter sync (DDP) or use RPC to invoke remote stages. Use activation checkpointing to reduce memory and support model/optimizer checkpointing per-rank.

Example using torch.distributed.pipeline.sync.Pipe (recommended when model is sequentially partitionable):

python

# assume torch.distributed.init_process_group has been called and ranks assigned
from torch.distributed.pipeline.sync import Pipe
import torch.nn as nn

# define full model layers list and split sizes based on ranks
layers = nn.Sequential(... large list ...)
# split_sizes could be determined by profiling FLOPs/memory per layer
partition = [layers[:L1], layers[L1:L2], layers[L2:]]
# each rank constructs only its partition and you create a Pipe on a single process group
model = Pipe(nn.Sequential(*partition_for_this_rank), chunks=NUM_MICROBATCHES)
# wrap each stage with DDP for parameter sync across data-parallel replicas
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_gpu])

If topology is non-sequential or you need custom RPC scheduling, use torch.distributed.rpc to define remote modules and implement a driver that:- Splits each mini-batch into micro-batches (m)- Issues forward RPCs to stage0 for micro-batch i, pipelined so stage k receives micro-batch i with appropriate delay- Maintains buffers of activation futures; once all forwards complete, trigger backward RPCs in reverse order (1F1B): while new forwards continue, interleave backward passes to free activations- Use torch.autograd.backward remotely or send gradients back; prefer torch.distributed.autograd with rpc.rpc_sync + torch.distributed.optim for end-to-end autograd across RPC.

Key engineering details:- Partitioning: profile layer memory/FLOPs; aim for balanced compute so no stage is a bottleneck. Consider grouping expensive ops (attention, large MLP) and splitting within layers only if you can shard tensors.- Micro-batches: choose chunks m so pipeline latency ~ (#stages + m -1)*micro_latency; throughput increases with m until bottleneck stage saturates. Typical m between 4–16. Use 1F1B schedule to reduce memory vs naive accumulation.- Parameter sync: combine pipeline per-stage with DDP across data-parallel replicas. Run DDP inside each stage so gradients are synchronized after backward; for RPC autograd, use torch.distributed.autograd and DistributedOptimizer to apply gradients across ranks.- Checkpointing: save per-rank state_dict (model, optimizer, RNG) frequently. For activation/compute checkpointing, use torch.utils.checkpoint.checkpoint to trade compute for memory. Use coordinated global checkpoint: rank0 triggers saving remote ranks via RPC to get consistent global step.- Fault tolerance: keep metadata (partition map, global step) in reliable storage; restore by reloading per-rank checkpoints.

Measuring latency vs throughput:- Latency per sample ≈ (pipeline_fill_time + compute_time_per_microbatch * stages) / m; measure end-to-end time per micro-batch and per-sample.- Throughput ≈ (m * batch_size) / steady_state_time_per_batch; run synthetic benchmarks varying m, micro-batch size, and number of pipeline stages.- Profiling: use NVProf/nsight or PyTorch profiler with RPC traces to identify stragglers. Track GPU utilization, interconnect bandwidth, and stage-wise compute time.- Trade-offs: larger m improves throughput but increases end-to-end latency and memory for activations; more stages reduce per-GPU memory but increase communication and latency; aggressive activation checkpointing reduces memory at cost of extra recompute.

Edge cases & best practices:- Avoid imbalanced stages; if unavoidable, replicate bottleneck stage (model-parallel intra-stage or tensor parallelism).- Align micro-batch sizes with optimizer batch accumulation to preserve effective batch size.- Use asynchronous communication carefully; prefer synchronous 1F1B to keep correctness simple.This design balances memory, compute, and communication to train very large models efficiently.

MediumTechnical

41 practiced

Write a TensorFlow Keras Callback (subclassing tf.keras.callbacks.Callback) that implements early stopping with a patience parameter and saves the best model weights to a provided path. It should monitor a named metric and optionally restore best weights at the end of training.

Sample Answer

Approach: Subclass tf.keras.callbacks.Callback to monitor a metric each epoch, track the best value, count patience, save best model weights to disk, and optionally restore best weights at training end.

python

import os
import numpy as np
import tensorflow as tf

class EarlyStoppingWithCheckpoint(tf.keras.callbacks.Callback):
    def __init__(self, monitor='val_loss', patience=5, mode='auto',
                 save_path='best_weights.h5', restore_best_weights=True, verbose=1):
        super().__init__()
        self.monitor = monitor
        self.patience = int(patience)
        self.save_path = save_path
        self.restore_best_weights = bool(restore_best_weights)
        self.verbose = verbose

        if mode not in ('auto', 'min', 'max'):
            raise ValueError("mode must be 'auto', 'min' or 'max'")
        if mode == 'min':
            self.monitor_op = np.less
            self.best = np.Inf
        elif mode == 'max':
            self.monitor_op = np.greater
            self.best = -np.Inf
        else:  # auto
            if 'loss' in self.monitor or self.monitor.startswith('val_'):
                self.monitor_op = np.less
                self.best = np.Inf
            else:
                self.monitor_op = np.greater
                self.best = -np.Inf

        self.wait = 0
        self.stopped_epoch = 0

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        current = logs.get(self.monitor)
        if current is None:
            if self.verbose:
                print(f"EarlyStoppingWithCheckpoint: Metric '{self.monitor}' is not available. Available: {list(logs.keys())}")
            return

        if self.monitor_op(current, self.best):
            self.best = current
            self.wait = 0
            # ensure directory exists
            os.makedirs(os.path.dirname(self.save_path) or '.', exist_ok=True)
            # save best weights
            self.model.save_weights(self.save_path)
            if self.verbose:
                print(f"Epoch {epoch+1}: {self.monitor} improved to {current:.5f}, saving model to {self.save_path}")
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch + 1
                self.model.stop_training = True
                if self.verbose:
                    print(f"Epoch {epoch+1}: early stopping triggered. No improvement in {self.patience} epochs.")

    def on_train_end(self, logs=None):
        if self.restore_best_weights and os.path.exists(self.save_path):
            try:
                self.model.load_weights(self.save_path)
                if self.verbose:
                    print(f"Restored best weights from {self.save_path}")
            except Exception as e:
                if self.verbose:
                    print(f"Failed to restore weights from {self.save_path}: {e}")

Key points:- Tracks metric and saves weights whenever improvement occurs.- Patience controls how many epochs without improvement to wait.- restore_best_weights loads the saved weights at training end.- Lightweight: O(1) work per epoch plus I/O when saving weights.Edge cases:- Metric not present in logs; path directory missing (handled); permission errors when writing; training restored without saved file.Alternatives:- Use built-in tf.keras.callbacks.EarlyStopping + ModelCheckpoint combination for production.

MediumTechnical

44 practiced

Explain how to save and restore complete training checkpoints in both TensorFlow and PyTorch, including model weights, optimizer state, learning rate scheduler state, epoch counters, and RNG states. Provide concise code snippets for saving/loading and explain why restoring optimizer state is necessary to resume training reliably.

Sample Answer

Brief approach: save everything needed to exactly resume training — model weights, optimizer state, LR scheduler state, current epoch/step counters, and random number generator (RNG) states for Python, NumPy and framework RNGs (PyTorch/TensorFlow). Restoring optimizer & scheduler ensures identical parameter momentum/accumulators and learning rate progression so training dynamics continue reliably.

PyTorch (recommended pattern)

python

# Save
checkpoint = {
  "epoch": epoch,
  "model_state": model.state_dict(),
  "optim_state": optimizer.state_dict(),
  "sched_state": scheduler.state_dict() if scheduler else None,
  "rng_python": importlib.reload(__import__("random")).getstate(),
  "rng_numpy": np.random.get_state(),
  "rng_torch": torch.get_rng_state(),
  "rng_cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
}
torch.save(checkpoint, "ckpt.pth")

# Load
ckpt = torch.load("ckpt.pth", map_location=device)
model.load_state_dict(ckpt["model_state"])
optimizer.load_state_dict(ckpt["optim_state"])
if scheduler and ckpt["sched_state"]: scheduler.load_state_dict(ckpt["sched_state"])
epoch = ckpt["epoch"]
random.setstate(ckpt["rng_python"])
np.random.set_state(ckpt["rng_numpy"])
torch.set_rng_state(ckpt["rng_torch"])
if ckpt["rng_cuda"]: torch.cuda.set_rng_state_all(ckpt["rng_cuda"])

TensorFlow / Keras (TF2)

python

# Save model + optimizer + other metadata
model.save_weights("model_weights")
# Save optimizer & scheduler via tf.train.Checkpoint
ckpt = tf.train.Checkpoint(step=tf.Variable(epoch), 
                           optimizer=optimizer,
                           model=model)
manager = tf.train.CheckpointManager(ckpt, "./tf_ckpts", max_to_keep=3)
manager.save()  # saves checkpoint-<N>

# Restore
ckpt = tf.train.Checkpoint(step=tf.Variable(0), optimizer=optimizer, model=model)
ckpt.restore(tf.train.latest_checkpoint("./tf_ckpts")).expect_partial()
epoch = int(ckpt.step.numpy())
# RNGs
import random, numpy as np
random.setstate(python_state)
np.random.set_state(numpy_state)
tf.random.set_seed(seed)  # TF RNG is typically seeded across graph

Why restore optimizer state: optimizers like Adam, RMSProp, and SGD with momentum maintain internal accumulators (momentum buffers, running averages, adaptive learning rate statistics). If you only restore weights but reinitialize the optimizer, those accumulators reset, causing a sudden change in update magnitudes and breaking reproducibility or hurting convergence — training may diverge or lose progress. Restoring scheduler state preserves LR schedule/step counts so learning rate resumes at the correct point.

Practical tips:- Save checkpoints atomically (write to temp then rename).- Version checkpoints if code/architecture changes.- Test full restore by running a few steps and comparing losses/gradients to pre-save run.

EasyTechnical

39 practiced

What is gradient clipping and why is it used? Explain the difference between clipping by norm and clipping by value. Provide a short PyTorch code snippet using torch.nn.utils.clip_grad_norm_ and discuss when gradient clipping is a bandaid vs a valid tool.

Sample Answer

Gradient clipping is a technique to limit the magnitude of gradients during backpropagation to prevent exploding gradients and stabilize training, especially in deep or recurrent networks. It does not change the model architecture or optimizer logic — it only scales or truncates gradients before the optimizer step.

Difference:- Clipping by value: clips each individual gradient element to a fixed range [−c, c]. It’s simple but can change gradient directions.- Clipping by norm: rescales the entire gradient vector for a parameter (or across parameters) if its L2 norm exceeds a threshold. It preserves gradient direction but reduces magnitude uniformly.

PyTorch snippet (clip by norm):

python

import torch
from torch import nn, optim
from torch.nn.utils import clip_grad_norm_

model = nn.Sequential(nn.Linear(100,50), nn.ReLU(), nn.Linear(50,1))
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# inside training loop after loss.backward()
max_norm = 1.0
clip_grad_norm_(model.parameters(), max_norm)  # rescales gradients if total norm > max_norm
optimizer.step()
optimizer.zero_grad()

When to use:- Valid tool: when training large/deep RNNs, transformers, or when occasional large gradients destabilize training; it reliably prevents NaNs and divergence.- Bandaid: if gradients explode due to bad model design, learning rate far too high, incorrect loss scaling, or data issues. Relying solely on clipping can hide underlying problems — investigate architecture, initialization, learning rate schedule, and loss formulation.

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.