10.1. What is a framework?
At its core, a machine learning framework has two main jobs: providing programming abstractions for users to define and train models, and making sure those abstractions execute efficiently on hardware. Let's look at the key building blocks a framework provides to the user, using PyTorch as an example.
First, a framework provides the Tensor data structure supporting different shapes and data types (dtype), as introduced in Tensor and Common data types, along with mathematical operators (ops) that transform these tensors, as introduced in Ops. The following snippet shows how PyTorch exposes these APIs to the user.
import torch
# Tensors with different shapes and data types on the GPU
token_ids = torch.tensor([101, 2054, 2003], dtype=torch.int64, device="cuda") # 1D: (3,)
x = torch.randn(2, 4, dtype=torch.bfloat16, device="cuda") # 2D: (2, 4)
w = torch.randn(4, 3, dtype=torch.bfloat16, device="cuda") # 2D: (4, 3)
b = torch.zeros(3, dtype=torch.bfloat16, device="cuda") # 1D: (3,)
batch = torch.ones(8, 16, 64, dtype=torch.float32, device="cuda") # 3D: (8, 16, 64)
# Applying common ops: matmul, addition, ReLU, and softmax
z = x @ w + b
a = torch.relu(z)
probs = torch.softmax(a, dim=-1)
Second, frameworks group tensors and ops into higher-level components called layers (or modules), such as linear and normalization layers, to make it easier to construct models, as shown in the snippet below.
import torch.nn as nn
# Built-in layers that package tensors and ops together
linear = nn.Linear(in_features=4, out_features=3, device="cuda", dtype=torch.bfloat16)
norm = nn.LayerNorm(normalized_shape=3, device="cuda", dtype=torch.bfloat16)
# Stacking layers to construct a model
model = nn.Sequential(
nn.Linear(4, 16),
nn.ReLU(),
nn.Linear(16, 3),
).to(device="cuda", dtype=torch.bfloat16)
output = model(x)
Third, frameworks automatically compute gradients for model parameters during training with a single backward() call on the loss, as shown in the snippet below.
# Forward pass through the linear layer and scalar loss computation
loss = linear(x).sum()
# Automatically compute gradients for the layer's parameters
loss.backward()
print(linear.weight.grad)
Beyond these three core abstractions, frameworks also provide utilities such as optimizers, data loaders, and loss functions. Modern frameworks also include modules for distributed execution and compilers, which we will explore in the following chapters.