GPU and Device

import torch

# Check availability
torch.cuda.is_available()         # True if CUDA GPU is accessible
torch.backends.mps.is_available() # True on Apple Silicon (Metal)
torch.cuda.device_count()         # number of GPUs
torch.cuda.current_device()       # index of current GPU
torch.cuda.get_device_name(0)     # GPU name string

# Universal device selection
device = (
    'cuda'  if torch.cuda.is_available()
    else 'mps' if torch.backends.mps.is_available()
    else 'cpu'
)
device = torch.device(device)

# Tensors
t = t.to(device)
t = t.to('cuda')
t = t.to('cuda:0')          # specific GPU index
t = t.cuda()                # shortcut
t = t.cpu()

# With options
t = t.to(device, dtype=torch.float16, non_blocking=True)

# Models
model = model.to(device)
model = model.cuda()

# Rule: always move model first, then data inside the loop
model = model.to(device)
for x, y in loader:
    x = x.to(device, non_blocking=True)   # non_blocking with pin_memory
    y = y.to(device, non_blocking=True)

# Pinned memory enables faster CPU→GPU transfers via DMA
t = torch.zeros(1000, 1000).pin_memory()
t = t.to(device, non_blocking=True)   # async transfer

# In DataLoader
loader = DataLoader(dataset, batch_size=32, pin_memory=True,
                    num_workers=4)     # combines pin_memory + non_blocking

# Explicitly select device
torch.cuda.set_device(1)

# Context manager — all operations in this block use cuda:1
with torch.cuda.device(1):
    model = MyModel().cuda()
    out = model(x.cuda())

# Create tensor on specific device
t = torch.empty(3, 3, device='cuda:1')
t = torch.empty(3, 3, device=torch.device('cuda', 1))

# Query
torch.cuda.memory_allocated()          # bytes in use (current tensors)
torch.cuda.max_memory_allocated()      # peak usage since last reset
torch.cuda.memory_reserved()           # bytes reserved by the caching allocator
torch.cuda.max_memory_reserved()       # peak reserved

# Human-readable summary
print(torch.cuda.memory_summary(device))

# Free unused cached memory (does not free in-use allocations)
torch.cuda.empty_cache()

# Reset peak stats
torch.cuda.reset_peak_memory_stats()

# Global memory
total, reserved, free = torch.cuda.mem_get_info()   # bytes

# 1. Use mixed precision
from torch.amp import autocast
with autocast('cuda', dtype=torch.float16):
    output = model(x)

# 2. Gradient checkpointing
from torch.utils.checkpoint import checkpoint
out = checkpoint(segment, x)

# 3. Delete intermediates and empty cache
del intermediate_tensor
torch.cuda.empty_cache()

# 4. Accumulate gradients over smaller batches
# 5. Use in-place operations where safe
# 6. Reduce batch size or input resolution

# Streams allow overlapping compute with memory transfers
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

with torch.cuda.stream(s1):
    y1 = model_a(x1)

with torch.cuda.stream(s2):
    y2 = model_b(x2)

torch.cuda.synchronize()    # wait for all streams on current device

# Current stream
current = torch.cuda.current_stream()
current.synchronize()

# CUDA events (for fine-grained timing)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
out = model(x)
end.record()
torch.cuda.synchronize()
elapsed_ms = start.elapsed_time(end)

# Always synchronize before timing!
torch.cuda.synchronize()
import time
t0 = time.perf_counter()
for _ in range(100):
    out = model(x)
torch.cuda.synchronize()
elapsed = (time.perf_counter() - t0) / 100

# CUDA event timing (more accurate)
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
out = model(x)
end.record()
torch.cuda.synchronize()
print(start.elapsed_time(end), 'ms')

from torch.amp import autocast, GradScaler

scaler = GradScaler('cuda')

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad(set_to_none=True)

    with autocast('cuda', dtype=torch.float16):
        output = model(x)
        loss = criterion(output, y)

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)
    scaler.update()

with autocast('cuda', dtype=torch.bfloat16):
    output = model(x)
    loss = criterion(output, y)
loss.backward()
optimizer.step()

# Benchmark mode: tries multiple algorithms and picks the fastest
# Good for: fixed input sizes, production training
torch.backends.cudnn.benchmark = True

# Deterministic mode: always use the same algorithm
# Good for: reproducibility
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False   # turn off benchmark when deterministic

# Disable cuDNN entirely (rarely needed)
torch.backends.cudnn.enabled = False

# Capture and replay a static computation graph
# Eliminates Python overhead — best for fixed-shape batch training

# Warmup
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):
        optimizer.zero_grad(set_to_none=True)
        y_pred = model(x)
        loss = criterion(y_pred, y)
        loss.backward()
        optimizer.step()
torch.cuda.current_stream().wait_stream(s)

# Capture
g = torch.cuda.CUDAGraph()
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
    static_output = model(static_x)
    static_loss = criterion(static_output, static_y)
    static_loss.backward()

# Replay (copy data into static buffers, replay graph)
static_x.copy_(real_x)
static_y.copy_(real_y)
g.replay()
optimizer.step()

device = torch.device('mps')

model = model.to(device)
x = x.to(device)

# Check MPS availability
torch.backends.mps.is_available()
torch.backends.mps.is_built()

# Fallback for unsupported ops
import os
os.environ['PYTORCH_ENABLE_MPS_FALLBACK'] = '1'

# DDP launch (from command line)
# torchrun --nproc_per_node=4 --nnodes=1 train.py

# Inside train.py
import torch.distributed as dist
dist.init_process_group(backend='nccl')   # nccl for GPU, gloo for CPU
local_rank = int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
model = DDP(model.to(local_rank), device_ids=[local_rank])
dist.destroy_process_group()