PyTorch: From First Tensor to Distributed Training

Scalar  (0D):  42                          shape: ()
Vector  (1D):  [1, 2, 3]                   shape: (3,)
Matrix  (2D):  [[1, 2], [3, 4]]            shape: (2, 2)
3D:            A stack of matrices          shape: (D0, D1, D2)
4D:            A batch of 3D tensors        shape: (batch, channels, height, width)

import torch

# From Python data
torch.tensor([1, 2, 3])                      # int64 by default
torch.tensor([1.0, 2.0, 3.0])                # float32 by default
torch.tensor([[1, 2], [3, 4]], dtype=torch.float32)

# Filled tensors
torch.zeros(3, 4)                             # all 0.0
torch.ones(2, 3)                              # all 1.0
torch.full((3, 3), 3.14)                      # all 3.14
torch.empty(5, 5)                             # uninitialized (garbage)
torch.eye(4)                                  # identity matrix

# Ranges
torch.arange(0, 10, 2)                        # [0, 2, 4, 6, 8]
torch.linspace(0, 1, steps=5)                 # [0, 0.25, 0.5, 0.75, 1.0]

# Random
torch.rand(3, 4)                              # uniform [0, 1)
torch.randn(3, 4)                             # normal(0, 1)
torch.randint(0, 10, (3, 4))                  # random ints [0, 10)

# From another tensor's properties
x = torch.randn(3, 4, device='cuda', dtype=torch.float16)
y = torch.zeros_like(x)                       # same shape, dtype, device
z = torch.empty_like(x)

t = torch.tensor([[1, 2, 3],
                   [4, 5, 6]])

# What Python shows you:     What's actually in memory:
# [[1, 2, 3],                [1, 2, 3, 4, 5, 6]
#  [4, 5, 6]]                 ← one flat array

t.shape          # torch.Size([2, 3])
t.stride()       # (3, 1)
#                   ↑  ↑
#                   |  └── to move one step in dim 1 (next column): skip 1 element
#                   └───── to move one step in dim 0 (next row): skip 3 elements

t.storage_offset()  # 0 — where this tensor starts in storage

element[i, j] = storage[offset + i * stride[0] + j * stride[1]]

Example: t[1, 2]
= storage[0 + 1*3 + 2*1]
= storage[5]
= 6  ✓

t = torch.arange(12)         # [0, 1, 2, ..., 11]
# Storage: [0,1,2,3,4,5,6,7,8,9,10,11]
# Shape: (12,)   Stride: (1,)

a = t.view(3, 4)
# Storage: [0,1,2,3,4,5,6,7,8,9,10,11]  ← SAME memory!
# Shape: (3, 4)   Stride: (4, 1)

b = t.view(2, 2, 3)
# Storage: [0,1,2,3,4,5,6,7,8,9,10,11]  ← STILL same memory!
# Shape: (2, 2, 3)   Stride: (6, 3, 1)

# All three tensors share the SAME storage
t[0] = 99
print(a[0, 0])    # 99 — same memory!

t = torch.tensor([[1, 2, 3],
                   [4, 5, 6]])
# Shape: (2, 3)   Stride: (3, 1)   Contiguous: YES
# Storage: [1, 2, 3, 4, 5, 6]

t2 = t.T   # transpose
# Shape: (3, 2)   Stride: (1, 3)   Contiguous: NO
# Storage: [1, 2, 3, 4, 5, 6]  ← SAME storage, just different strides!
#
# To read t2 row-by-row:
# t2[0,0] = storage[0*1 + 0*3] = storage[0] = 1
# t2[0,1] = storage[0*1 + 1*3] = storage[3] = 4
# t2[1,0] = storage[1*1 + 0*3] = storage[1] = 2
# t2[1,1] = storage[1*1 + 1*3] = storage[4] = 5
# t2[2,0] = storage[2*1 + 0*3] = storage[2] = 3
# t2[2,1] = storage[2*1 + 1*3] = storage[5] = 6
#
# Logical view of t2: [[1, 4],     Memory access pattern: 0, 3, 1, 4, 2, 5
#                       [2, 5],     ← jumping around! Not sequential!
#                       [3, 6]]

t2.is_contiguous()  # False — elements aren't sequential in memory

# The fix:
t2.contiguous()          # creates a NEW storage with sequential layout
# Now: Storage: [1, 4, 2, 5, 3, 6]   Stride: (2, 1)   Contiguous: YES

# In practice, use reshape() to avoid thinking about it:
t2.reshape(-1)           # always works, copies only if needed

t = torch.randn(2, 3, 4)

# Reshape
t.view(6, 4)              # (6, 4) — must be contiguous
t.reshape(6, 4)            # (6, 4) — always works
t.view(-1)                 # (24,)  — flatten
t.view(2, -1)              # (2, 12) — infer dim

# Transpose / Permute
m = torch.randn(3, 4)
m.T                        # (4, 3) — same as m.t()
t.permute(2, 0, 1)         # (4, 2, 3) — reorder ALL dims

# Squeeze / Unsqueeze
x = torch.randn(1, 3, 1, 4)
x.squeeze()                # (3, 4) — remove ALL size-1 dims
x.squeeze(0)               # (3, 1, 4) — remove specific
x.squeeze(2)               # (1, 3, 4) — remove specific

y = torch.randn(3, 4)
y.unsqueeze(0)             # (1, 3, 4) — add dim at position 0
y.unsqueeze(-1)            # (3, 4, 1) — add dim at end
y[None, :]                 # (1, 3, 4) — same as unsqueeze(0)
y[:, :, None]              # (3, 4, 1) — same as unsqueeze(-1)

# Flatten
t.flatten()                # (24,) — all dims
t.flatten(start_dim=1)     # (2, 12) — flatten dims 1 and 2

# Concat / Stack
a = torch.randn(3, 4)
b = torch.randn(3, 4)
torch.cat([a, b], dim=0)   # (6, 4) — join along EXISTING dim
torch.stack([a, b], dim=0)  # (2, 3, 4) — create NEW dim
torch.cat([a, b], dim=1)   # (3, 8)

# Expand / Repeat
t = torch.randn(1, 3)
t.expand(4, 3)             # (4, 3) — NO copy, broadcasts
t.expand(4, -1)            # (4, 3) — -1 means "keep"
t.repeat(4, 1)             # (4, 3) — COPIES data
t.repeat(4, 2)             # (4, 6) — copies and tiles

a = torch.randn(3, 4)
b = torch.randn(3, 4)

# Element-wise
a + b                       # addition
a * b                       # element-wise multiply (Hadamard)
a / b                       # division
a ** 2                      # power
torch.sqrt(a.abs())
torch.exp(a)
torch.log(a.abs())
torch.clamp(a, min=-1, max=1)   # clip values
torch.relu(a)                   # max(0, x)

# In-place operations (modify tensor directly, save memory)
a.add_(b)                   # a = a + b, but no new tensor allocated
a.mul_(2)                   # a = a * 2
a.zero_()                   # a = all zeros
a.fill_(3.14)               # a = all 3.14
# WARNING: in-place ops on tensors that require grad can cause errors

# Reductions
a.sum()                     # scalar — sum of all elements
a.sum(dim=0)                # (4,) — sum across rows
a.sum(dim=1, keepdim=True)  # (3, 1) — keep dimension
a.mean()                    # mean of all
a.mean(dim=-1)              # mean along last dim
a.std()                     # standard deviation
a.max()                     # global max (scalar)
a.max(dim=1)                # returns (values, indices) tuple
a.argmax(dim=1)             # indices only
a.min()
a.norm(p=2)                 # L2 norm
a.norm(p=2, dim=1)          # per-row L2 norm

# Matrix operations
x = torch.randn(3, 4)
y = torch.randn(4, 5)
z = x @ y                   # (3, 5) matrix multiply — PREFERRED syntax
z = torch.mm(x, y)          # same, explicit
z = torch.matmul(x, y)      # same, also handles batched

# Batch matrix multiply
bx = torch.randn(8, 3, 4)   # batch of 8 matrices
by = torch.randn(8, 4, 5)
bz = torch.bmm(bx, by)      # (8, 3, 5)
bz = bx @ by                # same — @ handles batched

# Einstein summation (powerful shorthand)
torch.einsum('ij,jk->ik', x, y)              # matmul
torch.einsum('bij,bjk->bik', bx, by)          # batch matmul
torch.einsum('ij->i', x)                      # sum along columns
torch.einsum('ii->', torch.eye(4))             # trace

t = torch.arange(20).reshape(4, 5)

# Basic slicing (returns VIEWS — shared memory)
t[0]                 # first row
t[1:3]               # rows 1, 2
t[:, 2]              # column 2
t[::2]               # every other row
t[:, -1]             # last column

# Boolean masking (returns COPY)
mask = t > 10
t[mask]               # 1D tensor of elements where mask is True
t[t % 3 == 0]         # elements divisible by 3

# Fancy indexing (returns COPY)
indices = torch.tensor([0, 2, 3])
t[indices]            # rows 0, 2, 3

# Scatter / Gather (important for embeddings, loss functions)
# gather: select elements along a dim using index tensor
src = torch.randn(3, 4)
idx = torch.tensor([[0, 1, 2, 3], [3, 2, 1, 0], [0, 0, 0, 0]])
torch.gather(src, dim=1, index=idx)   # picks src[i, idx[i,j]] for each (i,j)

# where
torch.where(t > 10, t, torch.zeros_like(t))  # keep if >10, else 0

x = torch.tensor(2.0, requires_grad=True)
y = torch.tensor(3.0, requires_grad=True)

# Forward pass builds the graph:
a = x * y          # a.grad_fn = <MulBackward0>
b = a + x          # b.grad_fn = <AddBackward0>
c = b ** 2          # c.grad_fn = <PowBackward0>

# The graph looks like:
#   x ──┬── MulBackward ── AddBackward ── PowBackward → c
#   y ──┘        │              │
#                └──────────────┘ (x feeds into both Mul and Add)

c.backward()

# What happens internally:
# 1. Start at c, gradient = 1.0 (dc/dc = 1)
# 2. PowBackward: dc/db = 2*b = 2*(a+x) = 2*(6+2) = 16
# 3. AddBackward: dc/da = 16, dc/dx_via_add = 16
# 4. MulBackward: dc/dx_via_mul = 16 * y = 48, dc/dy = 16 * x = 32
# 5. Total dc/dx = dc/dx_via_add + dc/dx_via_mul = 16 + 48 = 64

print(x.grad)   # tensor(64.)
print(y.grad)   # tensor(32.)

# The graph is DESTROYED after backward() (by default)
c.backward()
# c.backward()  ← would crash! Graph is freed.
# To keep it: c.backward(retain_graph=True)

# Gradients ACCUMULATE (they're summed, not replaced)
x.grad           # has a value from first backward
loss.backward()  # x.grad += new gradients
# That's why you need optimizer.zero_grad() every iteration

# Leaf tensors vs intermediate tensors
x = torch.randn(3, requires_grad=True)   # LEAF — user-created, grad stored
y = x * 2                                 # INTERMEDIATE — grad not stored by default
y.retain_grad()                            # explicitly request storage

# no_grad — for inference and parameter updates
with torch.no_grad():
    output = model(input)    # no graph built, saves memory
    param -= lr * param.grad # manual update without tracking

# inference_mode — even stricter, even faster
with torch.inference_mode():
    output = model(input)    # cannot use these tensors in any grad computation later

# detach — disconnect a tensor from the graph
z = some_tensor.detach()     # z shares data but has no grad_fn
# Common use: when you need a tensor's VALUE but not its gradient history
# Example: target computation in reinforcement learning

import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_classes):
        super().__init__()
        
        # PARAMETERS — learned, in state_dict, moved by .to()
        self.linear1 = nn.Linear(input_dim, hidden_dim)  # has .weight and .bias params
        self.linear2 = nn.Linear(hidden_dim, num_classes)
        
        # BUFFER — not learned, but in state_dict, moved by .to()
        # Example: a fixed positional encoding or a running count
        self.register_buffer('step_count', torch.zeros(1))
        
        # PLAIN ATTRIBUTE — not in state_dict, not moved
        self.hidden_dim = hidden_dim  # just a config value
    
    def forward(self, x):
        self.step_count += 1               # buffers update without grad
        x = torch.relu(self.linear1(x))
        return self.linear2(x)

model = MyModel(64, 128, 10)

# Iterate all parameters (recursive through submodules)
for name, param in model.named_parameters():
    print(f"{name}: shape={param.shape}, requires_grad={param.requires_grad}")
# linear1.weight: shape=torch.Size([128, 64]), requires_grad=True
# linear1.bias: shape=torch.Size([128]), requires_grad=True
# linear2.weight: shape=torch.Size([10, 128]), requires_grad=True
# linear2.bias: shape=torch.Size([10]), requires_grad=True

# Iterate buffers
for name, buf in model.named_buffers():
    print(f"{name}: {buf}")
# step_count: tensor([0.])

# Iterate submodules
for name, module in model.named_modules():
    print(f"{name}: {type(module).__name__}")
# '': MyModel
# linear1: Linear
# linear2: Linear

# Count parameters
total_params = sum(p.numel() for p in model.parameters())
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)

# Linear (fully connected)
nn.Linear(in_features=64, out_features=128, bias=True)
# Applies: y = x @ W.T + b   where W is (128, 64), b is (128,)

# Convolution
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1)
# Input: (N, 3, H, W) → Output: (N, 64, H, W)

# Normalization
nn.BatchNorm1d(num_features=128)     # for (N, 128) or (N, 128, L)
nn.BatchNorm2d(num_features=64)      # for (N, 64, H, W)
nn.LayerNorm(normalized_shape=128)   # normalizes over last dim(s)
nn.GroupNorm(num_groups=8, num_channels=64)

# Activation
nn.ReLU()
nn.GELU()
nn.SiLU()               # also called Swish

# Dropout
nn.Dropout(p=0.1)        # randomly zeros elements during training

# Embedding
nn.Embedding(num_embeddings=10000, embedding_dim=256)
# Input: (N, L) of int indices → Output: (N, L, 256)

# Transformer building blocks
nn.MultiheadAttention(embed_dim=256, num_heads=8)
nn.TransformerEncoderLayer(d_model=256, nhead=8, dim_feedforward=1024)

# Sequential — chain layers linearly
model = nn.Sequential(
    nn.Linear(64, 128),
    nn.ReLU(),
    nn.Linear(128, 10),
)

# Forward hook — runs after a module's forward()
activations = {}
def save_activation(name):
    def hook(module, input, output):
        activations[name] = output.detach()
    return hook

model.linear1.register_forward_hook(save_activation('linear1'))
output = model(x)
print(activations['linear1'].shape)   # captured intermediate output

# Backward hook — runs when gradients flow through a module
def print_grad(name):
    def hook(module, grad_input, grad_output):
        print(f"{name} grad_output norm: {grad_output[0].norm():.4f}")
    return hook

model.linear2.register_full_backward_hook(print_grad('linear2'))
loss.backward()  # prints gradient norm as it flows through linear2

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# ── Setup ──
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MyModel(input_dim=64, hidden_dim=256, num_classes=10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, 
                          num_workers=4, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False,
                        num_workers=4, pin_memory=True)

# ── Training ──
for epoch in range(100):
    
    # ── Train phase ──
    model.train()  # enables dropout, batchnorm uses batch stats
    train_loss = 0.0
    
    for batch_idx, (data, target) in enumerate(train_loader):
        data = data.to(device, non_blocking=True)     # [1] move to GPU
        target = target.to(device, non_blocking=True)
        
        optimizer.zero_grad()       # [2] clear old gradients
        # Alternative: optimizer.zero_grad(set_to_none=True)  ← slightly faster
        
        output = model(data)        # [3] forward pass — builds computation graph
        loss = criterion(output, target)  # [4] compute loss
        
        loss.backward()             # [5] backward pass — compute gradients
        
        # [6] Optional: gradient clipping (prevents exploding gradients)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()            # [7] update parameters using gradients
        
        train_loss += loss.item()   # .item() extracts Python float, frees tensor
    
    # ── Validation phase ──
    model.eval()   # disables dropout, batchnorm uses running stats
    val_loss = 0.0
    correct = 0
    
    with torch.no_grad():  # [8] no gradient computation — saves memory + speed
        for data, target in val_loader:
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            
            output = model(data)
            val_loss += criterion(output, target).item()
            
            pred = output.argmax(dim=1)
            correct += (pred == target).sum().item()
    
    # ── Scheduler step ──
    scheduler.step()   # [9] update learning rate
    
    # ── Logging ──
    avg_train = train_loss / len(train_loader)
    avg_val = val_loss / len(val_loader)
    accuracy = correct / len(val_dataset)
    lr = optimizer.param_groups[0]['lr']
    print(f"Epoch {epoch} | Train: {avg_train:.4f} | Val: {avg_val:.4f} | "
          f"Acc: {accuracy:.4f} | LR: {lr:.6f}")

# SGD — simple, needs momentum for good results
torch.optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=1e-4)

# Adam — adaptive learning rate per parameter, default choice
torch.optim.Adam(params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8)

# AdamW — Adam with decoupled weight decay (PREFERRED for transformers)
torch.optim.AdamW(params, lr=1e-3, weight_decay=0.01)
# AdamW applies weight decay DIRECTLY to parameters, not through gradients
# This is mathematically different from Adam + L2 regularization

optimizer = torch.optim.AdamW([
    {'params': model.backbone.parameters(), 'lr': 1e-5},   # low LR for pretrained
    {'params': model.head.parameters(), 'lr': 1e-3},       # high LR for new layers
], weight_decay=0.01)

# Step decay — multiply LR by gamma every step_size epochs
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# Cosine annealing — smooth decay from initial LR to near zero
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# Cosine with warm restarts — periodic resets
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
    optimizer, T_0=50, T_mult=2, eta_min=1e-6
)

# Reduce on plateau — lower LR when metric stops improving
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='min', factor=0.5, patience=10
)
# Note: this one takes a metric: scheduler.step(val_loss)

# Linear warmup (manual — PyTorch doesn't have a built-in one)
def warmup_lambda(current_step):
    warmup_steps = 1000
    if current_step < warmup_steps:
        return current_step / warmup_steps
    return 1.0
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_lambda)

# Warmup + cosine (common for transformer training)
def warmup_cosine(step):
    warmup_steps = 1000
    total_steps = 50000
    if step < warmup_steps:
        return step / warmup_steps
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    return 0.5 * (1 + math.cos(math.pi * progress))

scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warmup_cosine)

model = nn.Linear(3, 2)
print(model.state_dict())
# OrderedDict([
#   ('weight', tensor([[ 0.12, -0.34,  0.56], [-0.78, 0.90, -0.12]])),
#   ('bias',   tensor([0.23, -0.45]))
# ])
# Keys = parameter names (hierarchical, dot-separated for nested modules)
# Values = tensor data

# PATTERN 1: Save only state dict (RECOMMENDED)
torch.save(model.state_dict(), 'model_weights.pt')

# Load:
model = MyModel(...)       # must create model with same architecture first
model.load_state_dict(torch.load('model_weights.pt', weights_only=True))

# PATTERN 2: Save full checkpoint (for resuming training)
torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'loss': loss,
    'config': config,
}, 'checkpoint.pt')

# Load:
ckpt = torch.load('checkpoint.pt', weights_only=False)
model.load_state_dict(ckpt['model_state_dict'])
optimizer.load_state_dict(ckpt['optimizer_state_dict'])
start_epoch = ckpt['epoch'] + 1

# PATTERN 3: Partial loading (fine-tuning, architecture changes)
state = torch.load('pretrained.pt', weights_only=True)
model.load_state_dict(state, strict=False)
# strict=False: ignores missing/unexpected keys
# Returns: <IncompatibleKeys(missing_keys=[...], unexpected_keys=[...])>

# Check availability
torch.cuda.is_available()          # True/False
torch.cuda.device_count()         # number of GPUs
torch.cuda.current_device()       # current GPU index
torch.cuda.get_device_name(0)     # "NVIDIA A100-SXM4-80GB"

# Memory info
torch.cuda.memory_allocated()      # currently used by tensors (bytes)
torch.cuda.memory_reserved()       # total reserved by caching allocator
torch.cuda.max_memory_allocated()  # peak usage (for debugging OOM)

# Device object
device = torch.device('cuda:0')    # specific GPU
device = torch.device('cuda')      # current default GPU

# Moving data
t = torch.randn(3, 4)
t_gpu = t.to(device)               # copy to GPU
t_gpu = t.cuda()                   # shorthand
t_cpu = t_gpu.cpu()                # back to CPU

# Create on GPU directly
t = torch.randn(3, 4, device=device)  # no transfer needed

# CRITICAL: operations require same device
a = torch.randn(3, device='cuda:0')
b = torch.randn(3, device='cuda:1')
# a + b  ← RuntimeError! Different devices.
# a + b.to('cuda:0')  ← works

# Emptying cache (doesn't free memory, just returns it to allocator)
torch.cuda.empty_cache()

# Synchronization
torch.cuda.synchronize()   # wait for all GPU operations to finish
# GPU ops are async — Python returns immediately, GPU runs in background
# synchronize() blocks until all GPU work is done
# Needed before timing GPU operations

from torch.amp import autocast, GradScaler

# With float16 (needs scaler)
scaler = GradScaler()
for data, target in loader:
    optimizer.zero_grad()
    with autocast(device_type='cuda', dtype=torch.float16):
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)         # internally unscales, then steps
    scaler.update()                # adjusts scale factor dynamically

# With bfloat16 (no scaler needed)
for data, target in loader:
    optimizer.zero_grad()
    with autocast(device_type='cuda', dtype=torch.bfloat16):
        output = model(data)
        loss = criterion(output, target)
    loss.backward()
    optimizer.step()

Step 1-3 (Reduce-Scatter): Each GPU sends one chunk to its neighbor.
After receiving, it ADDS the received chunk to its own.
After 3 steps, each GPU has the SUM of one chunk across all GPUs.

Step 4-6 (All-Gather): Each GPU sends its completed sum to neighbors.
After 3 steps, every GPU has the complete summed gradient.

Total data sent per GPU: 2 × (N/P) × (P-1)
  where N = gradient size, P = number of GPUs
Key insight: bandwidth per GPU is CONSTANT regardless of P.

import os
import torch
import torch.nn as nn
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader, DistributedSampler

def main():
    # ── Initialize distributed ──
    # torchrun sets RANK, LOCAL_RANK, WORLD_SIZE as env vars
    dist.init_process_group(backend="nccl")  # NCCL = fastest for GPU
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
    
    # ── Model ──
    model = MyModel().to(device)
    model = DDP(model, device_ids=[local_rank])
    
    # ── Data ──
    dataset = MyDataset("/data/train")
    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank, shuffle=True)
    loader = DataLoader(
        dataset, batch_size=32, sampler=sampler,
        num_workers=4, pin_memory=True, drop_last=True,
    )
    
    # ── Optimizer ──
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
    scaler = torch.amp.GradScaler()
    
    # ── Resume from checkpoint ──
    start_epoch = 0
    ckpt_path = "latest_checkpoint.pt"
    if os.path.exists(ckpt_path):
        # All ranks load (they all need the same state)
        ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
        model.module.load_state_dict(ckpt["model"])
        optimizer.load_state_dict(ckpt["optimizer"])
        start_epoch = ckpt["epoch"] + 1
        if rank == 0:
            print(f"Resumed from epoch {start_epoch}")
    
    # ── Training loop ──
    for epoch in range(start_epoch, 100):
        sampler.set_epoch(epoch)       # CRITICAL for proper shuffling
        model.train()
        
        for step, (data, target) in enumerate(loader):
            data = data.to(device, non_blocking=True)
            target = target.to(device, non_blocking=True)
            
            optimizer.zero_grad(set_to_none=True)
            
            with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
                output = model(data)
                loss = nn.functional.cross_entropy(output, target)
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            
            if rank == 0 and step % 100 == 0:
                print(f"Epoch {epoch} Step {step} Loss {loss.item():.4f}")
        
        # ── Save checkpoint (rank 0 only) ──
        if rank == 0:
            torch.save({
                "epoch": epoch,
                "model": model.module.state_dict(),   # .module!
                "optimizer": optimizer.state_dict(),
            }, ckpt_path)
        
        # Ensure all ranks wait for rank 0 to finish saving
        dist.barrier()
    
    dist.destroy_process_group()

if __name__ == "__main__":
    main()

# Launch: single node, 4 GPUs
torchrun --standalone --nproc_per_node=4 train.py

# Launch: 2 nodes, 4 GPUs each
# On node 0:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=0 \
         --master_addr=10.0.0.1 --master_port=29500 train.py
# On node 1:
torchrun --nnodes=2 --nproc_per_node=4 --node_rank=1 \
         --master_addr=10.0.0.1 --master_port=29500 train.py

1. sampler.set_epoch(epoch)
   Without this, every epoch uses the same data ordering.
   The sampler seeds its shuffle with the epoch number.

2. model.module.state_dict()
   DDP wraps your model. The real model is inside .module.
   Save model.module.state_dict(), not model.state_dict().
   (model.state_dict() adds "module." prefix to all keys)

3. Save on rank 0 only
   All ranks have identical parameters. Multiple saves = wasted I/O.

4. dist.barrier() after saving
   Prevents other ranks from trying to load a half-written checkpoint.

5. Don't use model.eval() with DDP and SyncBatchNorm
   SyncBatchNorm requires all ranks to participate. If one rank skips
   the forward pass (e.g., uneven data), it deadlocks.

6. drop_last=True
   If the last batch has fewer samples, different ranks may have
   different batch sizes, causing hanging during gradient sync.

7. map_location when loading checkpoints
   torch.load('ckpt.pt', map_location=device) — ensures tensors
   are loaded to the correct GPU, not all to GPU 0.

accumulation_steps = 4
# effective_batch = batch_size × accumulation_steps × world_size

for step, (data, target) in enumerate(loader):
    data, target = data.to(device), target.to(device)
    
    with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
        output = model(data)
        loss = nn.functional.cross_entropy(output, target)
        loss = loss / accumulation_steps
    
    # Skip gradient sync on non-boundary steps (saves communication)
    if (step + 1) % accumulation_steps != 0:
        with model.no_sync():      # DDP won't all-reduce here
            loss.backward()
    else:
        loss.backward()            # all-reduce happens here
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()

DDP (4 GPUs):
  GPU 0: [full model] [full optimizer state] [full gradients]
  GPU 1: [full model] [full optimizer state] [full gradients]
  GPU 2: [full model] [full optimizer state] [full gradients]
  GPU 3: [full model] [full optimizer state] [full gradients]
  Total memory per GPU: ~3x model size

FSDP (4 GPUs):
  GPU 0: [1/4 model] [1/4 optimizer] [1/4 gradients]
  GPU 1: [1/4 model] [1/4 optimizer] [1/4 gradients]
  GPU 2: [1/4 model] [1/4 optimizer] [1/4 gradients]
  GPU 3: [1/4 model] [1/4 optimizer] [1/4 gradients]
  Total memory per GPU: ~3/4 x model size

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

# Wrap policy: tells FSDP where to shard (usually per transformer block)
wrap_policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)

# Mixed precision config
mp_policy = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    auto_wrap_policy=wrap_policy,
    mixed_precision=mp_policy,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    device_id=local_rank,
)

# Training loop is the same as DDP
# BUT: model.state_dict() returns sharded tensors
# Use distributed checkpoint (DCP) for saving — see next section

from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    apply_activation_checkpointing,
)

# Checkpoint every transformer block — recompute activations during backward
apply_activation_checkpointing(
    model,
    check_fn=lambda module: isinstance(module, TransformerBlock),
)
# Tradeoff: ~30% slower but uses ~60% less memory for activations

import torch.distributed.checkpoint as dcp

# ── Save (ALL ranks participate) ──
state = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "epoch": epoch,
}
dcp.save(state, checkpoint_id="/checkpoints/step_10000")
# Each rank writes its shard to a separate file:
# /checkpoints/step_10000/.metadata
# /checkpoints/step_10000/__0_0.distcp
# /checkpoints/step_10000/__1_0.distcp
# ...

# ── Load (ALL ranks participate, reshards automatically) ──
state = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
dcp.load(state, checkpoint_id="/checkpoints/step_10000")
model.load_state_dict(state["model"])
optimizer.load_state_dict(state["optimizer"])

# KEY FEATURE: Save on 8 GPUs, load on 4 GPUs. DCP reshards automatically.

# ── NCCL Timeout (most common distributed failure) ──
# Symptom: training hangs, eventually NCCL timeout error
# Causes:
#   - One rank crashed or is stuck (check all processes)
#   - Uneven data across ranks (use drop_last=True)
#   - One rank hit an exception before a collective operation
# Debug:
os.environ["NCCL_DEBUG"] = "INFO"      # verbose NCCL logging
os.environ["NCCL_TIMEOUT"] = "1800"    # 30 min timeout (default is ~10 min)
# Also: TORCH_DISTRIBUTED_DEBUG=DETAIL

# ── OOM (Out of Memory) ──
print(torch.cuda.memory_summary())     # detailed breakdown
torch.cuda.max_memory_allocated() / 1e9  # peak GB

# ── Gradient anomalies ──
torch.autograd.set_detect_anomaly(True)  # catches NaN/Inf in backward
# Expensive! Use only for debugging.

# ── Finding unused parameters in DDP ──
model = DDP(model, find_unused_parameters=True)
# If any parameter doesn't contribute to loss, DDP will deadlock
# because it expects gradients from all parameters for all-reduce.
# find_unused_parameters=True handles this but is slower.

# PyTorch Profiler
from torch.profiler import profile, record_function, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    with_stack=True,
) as prof:
    with record_function("training_step"):
        output = model(data)
        loss = criterion(output, target)
        loss.backward()

print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))

# Check if data loading is the bottleneck
import time
for i, batch in enumerate(loader):
    if i == 0:
        t_start = time.perf_counter()
    if i == 10:
        per_batch = (time.perf_counter() - t_start) / 10
        print(f"Data loading: {per_batch*1000:.1f}ms per batch")
        break