PyTorch & TensorFlow: The Training Frameworks

You know what a neural network does: numbers forward, a loss, gradients backward, weights nudged. A dot on a whiteboard.

But a dot cannot run on hardware. Somewhere, that math must become code — code that:

• moves billions of numbers through a GPU quickly,
• computes gradients automatically (backprop, implemented for you),
• survives a crash at step 40,000 of 50,000,
• and can be handed to a different machine later for serving.

That "somewhere" is the training framework. And in 2024–2026 the choice is mostly: PyTorch vs TensorFlow, with JAX as the third pole.

So the question becomes

It changes a lot while cooking: debugging, experimentation speed, who on your team already knows it, which libraries exist. It changes surprisingly little after plating: serving systems consume the exported dish, not the stove. Holding those two halves in your head — training-time concerns vs artifact boundaries — is the entire topic. Let us take them in order, starting with the one real philosophical difference between the two kitchens.

The two execution models, in one line each:

PyTorch (Meta, 2016) popularized eager execution:

• y = x @ W runs right now; you can print(y), inspect it in a debugger, set a breakpoint.
• Control flow is just Python: if/for over tensor shapes, variable-length inputs — trivial.
• The cost: each op launches separately, small kernels, Python overhead, no whole-program view to optimize across.

TensorFlow (Google, 2015) originally shipped static graphs: TF1 built a whole graph of ops first, then ran it in a session. Fast, portable, deployable — and miserable to debug ("why is my tensor NaN?" you could not just print; the value did not exist until the graph ran).

TF2 + Keras (2019+) conceded the ergonomics war: eager by default, with tf.function letting you wrap code in graph tracing when you want speed. Both kitchens now offer both modes.

PyTorch won the research race — the majority of NeurIPS/ICML reference code and virtually all LLM training (Hugging Face Transformers, torcheval, FSDP) is PyTorch-first. TensorFlow stays strong where its early production bet matured: TF Serving pipelines, TPU access via tf.distribute, and legacy enterprise stacks. Tellingly, Keras 3 (2023) became multi-backend — it runs on PyTorch, JAX, or TensorFlow — signaling where the ecosystem gravity sits.

Same math — one gradient-descent step — cooked two ways.

Eager (PyTorch). Every line is the computation, happening now:

code
Copy
 loss = ((model(x) - y)**2).mean()     # runs immediately: 0.8471...
 loss.backward()                        # gradients land in w.grad
 print(w.grad)                          # tensor([1.7, -0.3])  <- you SEE it
 w = w - lr * w.grad                    # tensor step, plain Python

You can print anything at any line. That is the whole drug.

Graph (TF1 style, and mentally what tf.function/XLA do): you first describe the recipe — loss → grads → update — as a symbolic program over placeholders. Nothing runs. The compiler then reads the whole recipe and rewrites it: fuse "subtract, square, mean" into one kernel instead of three round-trips to memory, keep intermediates in registers, plan the layout. Then execute, often much faster — but your print no longer appears where you wrote it; you inspect via special hooks.

The numeric payoff is real. On the same model with no code changes:

• plain eager step: every tiny op launched separately, Python between them
• compiled step (torch.compile or XLA): fused kernels → typically 1.3–2× faster training and inference

Which style "wins" depends on what you are doing this quarter: searching for an architecture (eager, today) or grinding tokens through a known one at scale (compiled, and you will meet the compilers next section).

Where the framework actually sits in the ML pipeline:

code
Copy
   TRAINING WORLD (framework matters here)          SERVING WORLD (artifact matters)
 ┌───────────────────────────────────┐            ┌──────────────────────────────┐
 │ code + gradients + GPUs           │            │                              │
 │ PyTorch / TensorFlow / JAX        │── EXPORT ──► ONNX · SavedModel ·         │
 │ trackers, launchers, checkpoints  │  artifacts  │ TorchScript · torch.export · │
 └───────────────────────────────────┘            │ safetensors                  │
                                                  │ → Triton / vLLM / TF Serving │
                                                  └──────────────────────────────┘

Read the arrow: the framework is a training-time concern. Serving systems consume exported artifacts — graphs and weights in exchange formats — so framework lock-in is weaker than it looks.

But note the fine print: the tooling around each framework (experiment trackers, model registries, exporters, distributed launchers) is integrated per-framework, so the choice ripples sideways through your MLOps even when it stops at the serving boundary. Sections 6–8 cover exactly where each framework is strong and how the pipeline gets built.

Carry this picture: you are choosing a kitchen for a restaurant.

• PyTorch is the tasting kitchen: burner at hand, you adjust the sauce spoon by spoon, and if it is wrong you know immediately which line of the recipe broke. Chefs (researchers) love it because experiments are cheap and mistakes are visible.
• TensorFlow launched as the factory kitchen: write the recipe on a card (graph), let the engineers optimize the card (fuse steps, remove dead branches), mass-produce identically. Debugging a bad batch means reading machinery, not tasting a spoon.
• TF2 is the factory admitting chefs need burners too: eager stoves now standard, recipe-cards (`tf.function) for the runs that need them.
• The compiled engines (torch.compile, XLA) are recipe-card printers that work on the tasting stove: keep writing freely, press a button, get the factory-optimized version of the same recipe — with one catch: the printer takes seconds-to-minutes to warm up (compilation latency).
• The freezer truck is the artifact boundary. The diners (users hitting your endpoint) never see the kitchen. They get the frozen dish: ONNX, SavedModel, TorchScript, safetensors. A dish made in a PyTorch kitchen and a dish made in a TF kitchen can ride the identical truck to the identical serving line.
• The pantry is the ecosystem: Hugging Face's pantry is PyTorch-branded; Google Cloud TPUs are a TF/JAX pantry. You can move pantries, but everything in it has to move with you.

When you interview or design, ask which part of the meal you are actually discussing: cooking (framework), printing recipes (compilation), the truck (export format), or the dining room (serving). Most "PyTorch vs TensorFlow" arguments are people confusing those four.

Eager speed is not optimal speed, and modern frameworks recover graph performance without rewriting models:

• torch.compile(model, mode="reduce-overhead") (PyTorch 2.0+, 2023): TorchDynamo traces your Python/ATen code → AOTAutograd extends the trace through the backward pass → TorchInductor emits fused Triton GPU kernels (or C++ CPU kernels). Typical 1.3–2× on training and inference with zero model changes. fullgraph=True plus CUDA graphs, and mode="max-autotune" (search kernel variants) squeeze more.
• XLA (TensorFlow: tf.function(jit_compile=True)): whole-graph compiler with operator fusion and memory-layout optimization. The default backend on TPUs; also available for GPUs.

Both fusion engines work on the same principle you met in the GPU topic: fewer round-trips to memory = fewer kernel launches = speed.

Both, however, suffer compilation latency: seconds to minutes for large graphs — a hidden cost in serverless cold starts and in notebooks whose shapes keep changing (recompile per shape). The 2025-era fixes are operational, not magical: bake compiled artifacts into your container images, and cache compiled graphs (FX graph caches, XLA service caches) in CI so the "printer warm-up" happens once, in the pipeline, not in front of a user.

import torch

model = MyTransformer()
compiled = torch.compile(model, mode="reduce-overhead")   # fast eager training/eval

# Freeze for a serving runtime that accepts TorchScript or Edge artifacts
example = torch.zeros(1, 128, dtype=torch.long)
exported = torch.export.export(model, (example,))          # stable graph, 2.2+
print(exported.graph_signature)

Production training runs on many GPUs, and each framework exposes a different primitive set:

PyTorch

• DDP: all-reduce data parallelism — each GPU gets a full model copy, gradients are averaged across GPUs every step.
• FSDP / FSDP2: parameters themselves are sharded across GPUs (the memory trick that trains models bigger than one card).
• Tensor/pipeline parallelism come largely via third-party libraries.
• Checkpoints are state_dict pickles — dangerous to load untrusted files (pickle executes arbitrary code). Use weights_only=True, or the safer, faster safetensors format.

TensorFlow

• tf.distribute.Strategy abstracts the topology: Mirrored (multi-GPU in one machine), MultiWorkerMirrored (many machines). One line change swaps the layout.
• SavedModel checkpoints are format-stable and safe to load.

Elasticity and failure semantics (both)

• Launchers: torchrun + torchx, or managed schedulers (Slurm, Ray Train) for gang scheduling — all-or-nothing worker placement.
• Recovery means a checkpoint every N steps, and you must budget the I/O: a 70B-parameter bf16 checkpoint is ~140 GB — write it slowly and your $2M cluster idles; write it too rarely and one GPU failure burns hours.

The design consequence, stated once so you keep it forever:

A registry entry without a documented, reshardable checkpoint format is a liability — you cannot move weights from 64 GPUs to 128, or feed a serving engine that expects a different tiling.

Decision heuristics for teams, in the order they should be weighed:

1. Research velocity / LLM workloads → PyTorch. The HF / TRL / vLLM ecosystem compounds — every new paper, dataset loader, and serving engine lands PyTorch-first.
2. TPU-backed large training → JAX or TensorFlow (XLA-native). If the accelerator is the constraint, the compiler stack is the framework.
3. Existing TF Serving fleets → stay on TF2 until migration pays off; the SavedModel export path still feeds production well, and "rewrite the working thing" is a cost, not a strategy.
4. Hiring reality → PyTorch proficiency now dominates the candidate pool.

And the escape hatch that makes the whole debate survivable: keep clean artifact boundaries. If your trackers, launchers, and registries talk to formats (safetensors, ONNX, SavedModel metadata) rather than frameworks, internal tooling stays agnostic and the truck can come from any kitchen.

In an interview, this is the mature answer shape: name the training-time differences honestly (eager ergonomics, ecosystem gravity, distribution primitives), then immediately draw the boundary line — serving consumes exported graphs, so framework choice and serving architecture are two decisions, not one. The people who conflate them are the ones stuck rewriting working models for fashion.