import torch
from torch.cuda.amp import autocast, GradScaler

model = Model().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
# GradScaler prevents FP16 underflow by scaling the loss up
# before backward, then scaling gradients back down before the step.
# Only needed for FP16; BF16 has enough dynamic range to skip this.
scaler = GradScaler()

for batch in dataloader:
    x, y = batch
    x, y = x.cuda(), y.cuda()
    
    optimizer.zero_grad()
    
    # Forward pass in mixed precision -- PyTorch automatically
    # casts operations to FP16 where safe and keeps FP32 where needed
    # (e.g., softmax, layer norm, loss computation)
    with autocast():
        output = model(x)
        loss = criterion(output, y)
    
    # Backward pass with loss scaling to prevent gradient underflow
    scaler.scale(loss).backward()
    scaler.step(optimizer)  # Unscales gradients, then steps
    scaler.update()         # Adjusts scale factor for next iteration

accumulation_steps = 4  # 4 mini-batches of 64 = effective batch of 256
effective_batch_size = batch_size * accumulation_steps

optimizer.zero_grad()
for i, batch in enumerate(dataloader):
    x, y = batch
    
    with autocast():
        output = model(x)
        # CRITICAL: divide loss by accumulation_steps so the accumulated
        # gradient is the same as if we had run the full batch at once.
        # Forgetting this division is a common bug that silently scales
        # your effective learning rate by accumulation_steps.
        loss = criterion(output, y) / accumulation_steps
    
    scaler.scale(loss).backward()  # Gradients accumulate in .grad
    
    # Only step the optimizer every accumulation_steps mini-batches
    if (i + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()  # Reset gradients for next accumulation cycle

from torch.utils.checkpoint import checkpoint

class CheckpointedBlock(nn.Module):
    """Wraps a block with gradient checkpointing.
    
    During forward: runs the block normally but does NOT save activations.
    During backward: re-runs the forward pass to recompute activations on the fly.
    """
    def __init__(self, block):
        super().__init__()
        self.block = block
    
    def forward(self, x):
        # use_reentrant=False is the recommended mode in modern PyTorch --
        # it handles edge cases with multiple autograd threads correctly.
        return checkpoint(self.block, x, use_reentrant=False)

# Apply to model -- typically checkpoint every Transformer block
model.layers = nn.ModuleList([
    CheckpointedBlock(layer) for layer in model.layers
])

model = nn.DataParallel(model)  # Simple but inefficient
output = model(input)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)

def train(rank, world_size):
    setup(rank, world_size)
    
    model = Model().to(rank)
    model = DDP(model, device_ids=[rank])
    
    # Training loop (each GPU gets different data)
    sampler = torch.utils.data.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, sampler=sampler, batch_size=batch_size)
    
    for epoch in range(epochs):
        sampler.set_epoch(epoch)  # Shuffle differently each epoch
        for batch in dataloader:
            # Normal training code
            ...

# Launch with:
# torchrun --nproc_per_node=4 train.py

# Split model into stages
class PipelinedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.stage1 = nn.Sequential(...).to('cuda:0')
        self.stage2 = nn.Sequential(...).to('cuda:1')
        self.stage3 = nn.Sequential(...).to('cuda:2')
    
    def forward(self, x):
        x = self.stage1(x.to('cuda:0'))
        x = self.stage2(x.to('cuda:1'))
        x = self.stage3(x.to('cuda:2'))
        return x

# Column-parallel linear layer
class ColumnParallelLinear(nn.Module):
    def __init__(self, in_features, out_features, world_size):
        super().__init__()
        assert out_features % world_size == 0
        self.local_out = out_features // world_size
        self.weight = nn.Parameter(torch.randn(self.local_out, in_features))
    
    def forward(self, x):
        local_out = F.linear(x, self.weight)
        # All-gather to combine outputs
        return all_gather(local_out)

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy

# Wrap model with FSDP
model = FSDP(
    model,
    auto_wrap_policy=transformer_auto_wrap_policy,
    mixed_precision=MixedPrecision(
        param_dtype=torch.bfloat16,
        reduce_dtype=torch.float32,
    ),
)

import deepspeed

model_engine, optimizer, _, _ = deepspeed.initialize(
    model=model,
    config={
        "train_batch_size": 256,
        "gradient_accumulation_steps": 4,
        "fp16": {"enabled": True},
        "zero_optimization": {
            "stage": 3,  # ZeRO Stage 3
            "offload_param": {"device": "cpu"},
            "offload_optimizer": {"device": "cpu"},
        },
    }
)

for batch in dataloader:
    loss = model_engine(batch)
    model_engine.backward(loss)
    model_engine.step()

from torch.utils.data import DataLoader

dataloader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,           # Parallel data loading -- use 4-8 workers per GPU.
                             # Too many workers wastes CPU/RAM; too few starves the GPU.
    pin_memory=True,         # Pre-allocates batches in page-locked (pinned) memory,
                             # which allows async CPU-to-GPU transfer (2-3x faster).
    prefetch_factor=2,       # Each worker prefetches 2 batches ahead.
    persistent_workers=True, # Keep worker processes alive between epochs --
                             # avoids the overhead of spawning new processes each epoch.
)