Autograd and Backpropagation in PyTorch

# io.thecodeforge: Standard Autograd and Backpropagation Flow
# This example walks through every step of the gradient computation
# so you can see what Autograd is doing at each stage.
import torch

# requires_grad=True marks x as a leaf tensor whose gradient we want.
# Autograd will track every operation performed on x.
x = torch.ones(2, 2, requires_grad=True)
print(f"x.requires_grad: {x.requires_grad}")
print(f"x.grad_fn: {x.grad_fn}")  # None — x is a leaf, not computed from another tensor

# Forward pass — each operation creates a Function node in the graph
y = x + 2
print(f"y.grad_fn: {y.grad_fn}")  # AddBackward0 — knows how to compute its gradient

z = y * y * 3
print(f"z.grad_fn: {z.grad_fn}")  # MulBackward0

out = z.mean()
print(f"out.grad_fn: {out.grad_fn}")  # MeanBackward0
print(f"Forward Pass Result: {out.item()}")

# Backward pass — traverse the graph in reverse, applying Chain Rule at each node
# PyTorch computes d(out)/dx automatically
out.backward()

# x.grad now contains d(out)/dx
# Manual derivation: out = z.mean() = (3 * (x+2)^2).mean()
# d(out)/dx = 3 * 2 * (x+2) / 4 = 3 * 2 * 3 / 4 = 4.5 at x=1
print(f"Gradients d(out)/dx:\n{x.grad}")
# Expected: all 4.5 because x was initialized to ones

# The graph is destroyed after backward() — calling again raises RuntimeError
# This is intentional memory management, not a limitation

# io.thecodeforge: Computational graph lifecycle and memory management
# Understanding this code is the key to diagnosing most PyTorch memory issues
import torch

# === GRAPH CONSTRUCTION ===
# Each operation creates a Function node that chains backward through next_functions
x = torch.randn(3, 3, requires_grad=True)
y = x * 2           # MulBackward0 node created, stores reference to x
z = y.sum()         # SumBackward0 node created, stores reference to y

print(f"z.grad_fn: {z.grad_fn}")  # SumBackward0
print(f"z.grad_fn.next_functions: {z.grad_fn.next_functions}")  # -> MulBackward0
print(f"MulBackward0.next_functions: {z.grad_fn.next_functions[0][0].next_functions}")
# -> AccumulateGrad (the leaf x)

# === RETENTION PHASE ===
# The graph is in memory between the forward pass completion and backward()
# Every intermediate tensor y is held alive by the graph during this window

# === DESTRUCTION — default behavior ===
z.backward()  # Graph traversed, gradients computed, graph freed
print(f"d(z)/dx:\n{x.grad}")  # tensor([[2., 2., 2.], ...])

# Attempting backward again raises RuntimeError — graph is gone
try:
    z.backward()
except RuntimeError as e:
    print(f"Expected error: {e[:80]}...")

# === retain_graph=True — use sparingly ===
x.grad = None  # Reset gradient before next backward
y2 = x * 3
z2 = y2.sum()

# retain_graph=True keeps the graph alive after backward
# Use ONLY when you need multiple backward passes on the same graph
# (e.g., computing gradients of multiple scalar outputs separately)
z2.backward(retain_graph=True)  # Graph retained
first_grad = x.grad.clone()
print(f"After first backward (retain_graph=True):\n{first_grad}")

z2.backward()  # Second backward works — but gradients are ACCUMULATED
print(f"After second backward (accumulated):\n{x.grad}")
# x.grad is now 2x the correct gradient — accumulated, not replaced
# This is why retain_graph in a loop causes exponential gradient growth

# === MEMORY LEAK PATTERN — what NOT to do ===
# This pattern causes GPU memory to grow linearly with training steps:
# loss_history = []
# for batch in dataloader:
#     loss = model(batch)
#     loss_history.append(loss)  # WRONG: holds entire graph for every batch
#     loss.backward()
#
# Correct:
# loss_history.append(loss.item())  # Python float, no graph reference

# io.thecodeforge: The correct training loop pattern — every decision annotated
import torch
import torch.nn as nn
import torch.optim as optim

model = nn.Linear(10, 1)
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.MSELoss()

# Dummy data
inputs = torch.randn(32, 10)   # batch of 32
targets = torch.randn(32, 1)

loss_history = []  # For monitoring — must store scalars, not tensors

for step in range(100):
    # STEP 1: Zero gradients BEFORE the forward pass.
    # Position matters: putting this after backward() instead of before
    # the next forward pass is equivalent for single-step loops but breaks
    # gradient accumulation workflows. Always put it before forward.
    optimizer.zero_grad()

    # STEP 2: Forward pass — builds the computational graph
    outputs = model(inputs)
    loss = criterion(outputs, targets)  # Scalar — averaged over batch by MSELoss

    # STEP 3: Log the scalar loss value WITHOUT a graph reference
    # loss is a tensor with requires_grad=True — storing it holds the entire graph
    # loss.item() returns a Python float — no graph reference, no memory leak
    loss_history.append(loss.item())  # Correct
    # loss_history.append(loss)       # WRONG — holds graph reference across steps

    # STEP 4: Backward pass — compute gradients for all parameters
    loss.backward()  # Graph is traversed and freed after this call

    # STEP 5: (Optional but recommended) Gradient clipping before optimizer step
    # Prevents a single bad batch from causing catastrophic weight updates
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

    # STEP 6: Update weights using computed gradients
    optimizer.step()

    # STEP 7: (Optional but valuable) Monitor gradient norms
    # A sudden spike here is the earliest signal of gradient issues
    if step % 10 == 0:
        grad_norm = sum(
            p.grad.norm().item() ** 2
            for p in model.parameters() if p.grad is not None
        ) ** 0.5
        print(f"Step {step:3d}: loss={loss_history[-1]:.4f}, grad_norm={grad_norm:.4f}")

print(f"\nFinal loss: {loss_history[-1]:.4f}")
print(f"All losses are Python floats (no graph refs): {type(loss_history[0])}")

# io.thecodeforge: Custom autograd Function examples
# Two real production use cases: numerically stable log-sum-exp
# and straight-through estimator for quantization-aware training
import torch
from torch.autograd import Function


class StableLogSumExp(Function):
    """
    Numerically stable log-sum-exp with a custom backward pass.

    Why this exists:
    The naive computation log(sum(exp(x))) overflows float32 when any
    element of x exceeds ~88. The stable version subtracts the maximum
    before exponentiating — mathematically identical but float32-safe.

    The gradient is softmax(x) * grad_output, computed in the stable form.
    """

    @staticmethod
    def forward(ctx, x, dim):
        ctx.dim = dim
        # Subtract max for numerical stability — does not change the result
        max_val = x.max(dim=dim, keepdim=True).values
        shifted = x - max_val
        exp_sum = shifted.exp().sum(dim=dim, keepdim=True)
        ctx.save_for_backward(x, max_val)
        return max_val.squeeze(dim) + exp_sum.log().squeeze(dim)

    @staticmethod
    def backward(ctx, grad_output):
        x, max_val = ctx.saved_tensors
        dim = ctx.dim
        # Stable softmax — same stability trick as forward
        shifted = (x - max_val).exp()
        softmax = shifted / shifted.sum(dim=dim, keepdim=True)
        # Chain rule: d(LSE)/dx_i = softmax_i
        # Multiply by upstream gradient and broadcast back to x shape
        return grad_output.unsqueeze(dim) * softmax, None  # None for dim argument


class StraightThroughEstimator(Function):
    """
    Straight-through estimator for quantization-aware training.

    Why this exists:
    Rounding (quantization) has zero gradient everywhere — no signal
    propagates backward through round() in standard Autograd.
    The straight-through estimator passes gradients straight through
    the quantization step as if it were an identity function.
    This allows gradient-based training of networks that use quantization
    at inference time.
    """

    @staticmethod
    def forward(ctx, x, num_bits=8):
        # Quantize to num_bits — this is the discrete rounding step
        scale = (2 ** num_bits - 1)
        return x.mul(scale).round().div(scale)

    @staticmethod
    def backward(ctx, grad_output):
        # Pass gradient through unchanged — ignore the discontinuity
        # This is the straight-through estimator approximation
        return grad_output, None  # None for num_bits


# === Verification with gradcheck ===
# gradcheck numerically estimates gradients via finite differences
# and compares against your custom backward() — always run this before deploying
print("Verifying StableLogSumExp gradients...")
x_check = torch.randn(4, 8, requires_grad=True, dtype=torch.float64)  # float64 for gradcheck precision
result = torch.autograd.gradcheck(
    lambda x: StableLogSumExp.apply(x, 1),
    (x_check,),
    eps=1e-6,
    atol=1e-4
)
print(f"StableLogSumExp gradcheck: {'PASSED' if result else 'FAILED'}")

# === Production usage ===
x = torch.randn(3, 10, requires_grad=True)
lse = StableLogSumExp.apply(x, dim=1)
loss = lse.sum()
loss.backward()
print(f"Custom gradient shape: {x.grad.shape}")
# Gradient sums to 1.0 per row because gradient of LSE is softmax
print(f"Gradient row sums (should be ~1.0): {x.grad.sum(dim=1)}")

# === Quantization example ===
weights = torch.randn(4, 4, requires_grad=True)
quantized = StraightThroughEstimator.apply(weights, 8)
loss_q = quantized.sum()
loss_q.backward()
print(f"\nQuantization gradient (straight-through, same shape as weights): {weights.grad.shape}")
print(f"Gradient values (all 1.0 — passed straight through): {weights.grad.unique()}")

# io.thecodeforge: Gradient checkpointing for memory-efficient large model training
# This pattern reduces peak GPU memory for a 24-layer transformer
# from ~40GB to ~12GB at the cost of ~30% more compute
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint


class TransformerBlock(nn.Module):
    """
    A standard transformer block: attention + FFN with residual connections.
    The expensive memory consumer: attention intermediate activations
    (sequence_len x sequence_len matrices) and FFN hidden states (4x expansion).
    """
    def __init__(self, d_model=512, nhead=8):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, nhead, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)

    def forward(self, x):
        # Without checkpointing: all of attn_out, ffn_out are stored for backward
        attn_out, _ = self.attn(x, x, x)
        x = self.norm1(x + attn_out)
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)
        return x


class CheckpointedTransformer(nn.Module):
    """
    Transformer with gradient checkpointing on expensive blocks.

    Memory profile comparison (12 layers, d_model=512, seq_len=512, batch=32):
    - Without checkpointing: ~24GB peak GPU memory
    - With checkpointing:    ~8GB peak GPU memory
    - Compute overhead:      ~30% more forward pass time
    """
    def __init__(self, n_layers=12, d_model=512):
        super().__init__()
        # Embedding is cheap to store — do not checkpoint it
        self.embedding = nn.Embedding(50000, d_model)
        self.layers = nn.ModuleList([
            TransformerBlock(d_model) for _ in range(n_layers)
        ])
        # Output projection is also cheap — do not checkpoint
        self.output_proj = nn.Linear(d_model, 50000)

    def forward(self, x):
        # Embedding — cheap, store the intermediate
        x = self.embedding(x)

        for layer in self.layers:
            # checkpoint() discards all intermediates inside TransformerBlock.forward()
            # During backward, it re-executes the block to reconstruct them
            # use_reentrant=False: use the newer, safer implementation (PyTorch 2.0+)
            # The old default (use_reentrant=True) has edge cases with double backward
            x = checkpoint(layer, x, use_reentrant=False)

        # Output projection — cheap, store the intermediate
        return self.output_proj(x)


# Demonstrate memory measurement
model = CheckpointedTransformer(n_layers=12, d_model=512)

# Use integer token IDs for the embedding layer
batch = torch.randint(0, 50000, (4, 64))  # batch_size=4, seq_len=64 tokens

torch.cuda.reset_peak_memory_stats() if torch.cuda.is_available() else None
out = model(batch)
loss = out.sum()
loss.backward()

if torch.cuda.is_available():
    peak_mb = torch.cuda.max_memory_allocated() / 1024**2
    print(f"Peak GPU memory: {peak_mb:.1f} MB")

print(f"Output shape: {out.shape}")
print(f"Embedding gradient exists: {model.embedding.weight.grad is not None}")
print(f"Layer 0 attention gradient exists: {model.layers[0].attn.in_proj_weight.grad is not None}")
print(f"All parameters received gradients: "
      f"{all(p.grad is not None for p in model.parameters())}")

# io.thecodeforge: Comparing inference contexts — performance and safety
import torch
import torch.nn as nn
import time

model = nn.Sequential(
    nn.Linear(1024, 2048),
    nn.ReLU(),
    nn.Linear(2048, 2048),
    nn.ReLU(),
    nn.Linear(2048, 1024),
    nn.ReLU(),
    nn.Linear(1024, 10)
)

x = torch.randn(256, 1024)
N_ITER = 1000

# === Benchmark 1: Default (graph tracking enabled) ===
# Builds the full computational graph for every forward pass
# Memory: graph + all intermediates allocated
start = time.perf_counter()
for _ in range(N_ITER):
    _ = model(x)
default_time = time.perf_counter() - start

# === Benchmark 2: torch.no_grad() ===
# Disables gradient computation but still increments version counters
# Appropriate for validation during training (safe to mix with tracked ops)
start = time.perf_counter()
for _ in range(N_ITER):
    with torch.no_grad():
        _ = model(x)
no_grad_time = time.perf_counter() - start

# === Benchmark 3: torch.inference_mode() ===
# Disables gradient computation AND version tracking
# Maximum performance for pure inference — use for production serving
start = time.perf_counter()
for _ in range(N_ITER):
    with torch.inference_mode():
        _ = model(x)
inference_time = time.perf_counter() - start

print(f"Default (graph tracking):   {default_time:.3f}s  (baseline)")
print(f"no_grad():                  {no_grad_time:.3f}s  ({(1 - no_grad_time/default_time)*100:.1f}% faster)")
print(f"inference_mode():           {inference_time:.3f}s  ({(1 - inference_time/default_time)*100:.1f}% faster)")

# === Correct usage patterns ===

# Pattern 1: Production inference — always use inference_mode
with torch.inference_mode():
    predictions = model(x)
    # predictions.requires_grad is False — cannot call backward on them
    # This is correct for serving: predictions are outputs, not training intermediates

# Pattern 2: Validation during training — no_grad is appropriate
# because you are still in a training context and may mix with gradient operations
loss = torch.tensor(0.5, requires_grad=True)  # example tracked tensor
with torch.no_grad():
    val_output = model(x)  # no gradient tracked for this
# loss.backward() still works here — no_grad only affected val_output

# Pattern 3: .detach() for logging a specific tensor
training_loss = torch.tensor(1.23, requires_grad=True)  # has grad_fn
loss_for_log = training_loss.detach()     # same value, no graph reference
loss_scalar = training_loss.item()        # Python float, no tensor at all

print(f"\nTraining loss requires_grad: {training_loss.requires_grad}")
print(f"Detached loss requires_grad: {loss_for_log.requires_grad}")
print(f".item() returns Python type: {type(loss_scalar)}")

# Pattern 4: .detach() to break gradient flow within a model
# Use when you want one branch to not affect gradients of shared layers
encoder_output = torch.randn(8, 128, requires_grad=True)
stop_gradient = encoder_output.detach()  # downstream ops see value but no gradient
print(f"\nStop-gradient tensor requires_grad: {stop_gradient.requires_grad}")