PyTorch and Memory

float32:

• fp32 is the default

Memory is determined by 1) the # of values in a tensor, 2) the data type of each value

x = torch.zeros(4, 8)
assert x.dtype == torch.float32 # default type
assert x.size() == torch.size([4, 8]) # shape of the tensor
assert x.numel() == 32 # total number of elements
assert x.element_size() == 4 # size of each element in bytes (float32 = 4 bytes)
assert get_memory_usage(x) == 4*8*4 == 128 # total memory in bytes (32 elements * 4 bytes each)

• bfloat16 uses the same memory as float16 but has the same dynamic range as float32
  • Resolution (part determined by fraction is worse)
• tldr: training with fp32 works, but requires lots of memory; training with fp8/fp16/bloat16 is riskier, can get instability
  • See mixed precision training

Moving things to the GPU:

Tensor Operations

def tensor_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Perform matrix multiplication on two tensors.
    """
    a = torch.ones(16, 32)
    b = torch.ones(32, 2)
 
    return a @ b
 

# jaxtyping
def jaxtyping_basics():
    x = torch.ones(2,2,1,3) # batch, seq, heads, hidden 
 
    # new jaxtyping way:
    x: Float[torch.Tensor, "batch seq heads hidden"] = torch.ones(2,2,1,3)
def einops_einsum():
    x: Float[torch.tensor, "batch seq1 hidden"] = torch.ones(2,3,4)
    y: Float[torch.tensor, "batch seq2 hidden"] = torch.ones(2,3,4)
 
    # old way 
    z = x @ y.transpose(-2, -1) # batch, seq1, seq2
 
    # new way
    z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")
 
def einops_reduce():
    x: Float[torch.tensor, "batch seq hidden"] = torch.ones(2,3,4)
 
    # old way
    y = x.mean(dim=1) # batch, hidden
 
    # new way
    y = reduce(x, "batch seq hidden -> batch hidden", "mean")
 
 

---

tensor_operations_flops

• A FLOP is a floating-point operation (e.g. +, *)
• FLOPs = total number of operations
• FLOP/s = FLOPs per second (hardware speed, also written FLOPS)

Real-world intuition

• GPT-3 (2020): ~3.14e23 FLOPs
• GPT-4 (2023 est.): ~2e25 FLOPs
• U.S. EO (revoked 2025): report any model trained with ≥1e26 FLOPs

GPU Specs (dense matmuls)

a100_flop_per_sec = 312e12  # A100, bfloat16 peak
h100_flop_per_sec = 1979e12 / 2  # H100, bfloat16 dense

Example: Linear Model

Let B = batch size, D = input dim, K = output dim

x = torch.ones(B, D)
w = torch.randn(D, K)
y = x @ w
actual_num_flops = 2 * B * D * K

• Matrix multiplication: 2 * m * n * p FLOPs
• Elementwise ops: m * n FLOPs
• Matrix addition: m * n FLOPs
• Matmuls dominate FLOPs in deep learning.

FLOP/s and MFU

actual_time = time_matmul(x, w)
actual_flop_per_sec = actual_num_flops / actual_time
promised_flop_per_sec = get_promised_flop_per_sec(device, x.dtype)
mfu = actual_flop_per_sec / promised_flop_per_sec

MFU ≥ 0.5 is considered good (model FLOPs utilization)

bfloat16 version

x = x.to(torch.bfloat16)
w = w.to(torch.bfloat16)
bf16_actual_time = time_matmul(x, w)
bf16_flop_per_sec = actual_num_flops / bf16_actual_time
bf16_promised = get_promised_flop_per_sec(device, x.dtype)
bf16_mfu = bf16_flop_per_sec / bf16_promised

---

gradients_basics

x = torch.tensor([1., 2, 3])
w = torch.tensor([1., 1, 1], requires_grad=True)
pred_y = x @ w
loss = 0.5 * (pred_y - 5).pow(2)
loss.backward()
# w.grad == tensor([1, 2, 3])

---

gradients_flops

Model:

x → h1 = x @ w1 → h2 = h1 @ w2 → loss = mean(h2^2)

num_forward_flops = 2*B*D*D + 2*B*D*K
 
# Backward pass:
# dL/dw2
num_backward_flops = 2 * B * D * K
 
# dL/dh1
num_backward_flops += 2 * B * D * K
 
# dL/dw1
num_backward_flops += 4 * B * D * D

→ Total: 6 * (# data points) * (# parameters) FLOPs

---

module_parameters

w = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))

• Xavier init (scales output ~N(0, 1))
• To avoid extreme values: use truncated normal:

nn.init.trunc_normal_(w, std=1/np.sqrt(input_dim), a=-3, b=3)

---

custom_model

class Linear(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
    def forward(self, x):
        return x @ self.weight
 
class Cruncher(nn.Module):
    def __init__(self, dim, num_layers):
        super().__init__()
        self.layers = nn.ModuleList([Linear(dim, dim) for _ in range(num_layers)])
        self.final = Linear(dim, 1)
    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x.squeeze(-1)

---

get_batch

def get_batch(data, batch_size, sequence_length, device):
    start_indices = torch.randint(len(data) - sequence_length, (batch_size,))
    x = torch.tensor([data[start:start + sequence_length] for start in start_indices])
    if torch.cuda.is_available():
        x = x.pin_memory()
    return x.to(device, non_blocking=True)

---

randomness

Set seeds for reproducibility:

torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)

---

data_loading

# Save
np.array([1,2,3]).tofile("data.npy")
 
# Load (lazy)
data = np.memmap("data.npy", dtype=np.int32)

---

optimizers

class AdaGrad(torch.optim.Optimizer):
    def step(self):
        for group in self.param_groups:
            lr = group["lr"]
            for p in group["params"]:
                grad = p.grad.data
                g2 = self.state[p].get("g2", torch.zeros_like(grad))
                g2 += grad**2
                self.state[p]["g2"] = g2
                p.data -= lr * grad / torch.sqrt(g2 + 1e-5)

• AdaGrad: accumulates squared gradients
• RMSProp: AdaGrad + exp average
• Adam: RMSProp + momentum

---

memory accounting

num_parameters = D * D * num_layers + D
num_activations = B * D * num_layers
num_gradients = num_parameters
num_optimizer_states = num_parameters
total_memory = 4 * (params + acts + grads + opt_states)

Assume float32 (4 bytes per value)

---

training loop

def train(...):
    for t in range(num_train_steps):
        x, y = get_batch(B)
        pred = model(x)
        loss = F.mse_loss(pred, y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

---

checkpointing

# Save
torch.save({
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}, "model_checkpoint.pt")
 
# Load
state = torch.load("model_checkpoint.pt")

---

mixed_precision_training

• Use lower precision (fp16, bf16, fp8) where possible
• Use float32 for weights, gradients

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    pred = model(x)

→ Supported with torch.cuda.amp or NVIDIA’s Transformer Engine for FP8

---