GPU and Hardware Considerations Questions

def estimate_gpu_memory_gb(
    P,
    bytes_per_activation_per_sample,
    batch_size,
    optimizer_multiplier=2.0,
    bytes_per_param=4,
    scratch_multiplier=0.1,
):
    """
    Estimate GPU memory usage.

    Args:
      P: int - number of model parameters
      bytes_per_activation_per_sample: float - bytes used by activations per sample
      batch_size: int
      optimizer_multiplier: float - e.g., Adam ~= 2 (moment estimates) or 1 for SGD
      bytes_per_param: int - bytes per parameter (4 for float32, 2 for float16)
      scratch_multiplier: float - fraction of (weights + activations) reserved for temporary buffers

    Returns:
      dict with:
        total_gb: float
        breakdown_bytes: dict of component sizes in bytes
        breakdown_gb: dict of component sizes in GB
    """
    if P < 0 or batch_size < 0 or bytes_per_activation_per_sample < 0:
        raise ValueError("P, batch_size and bytes_per_activation_per_sample must be non-negative")

    # Components in bytes
    weights = int(P * bytes_per_param)
    gradients = int(weights)  # typically one gradient tensor same dtype as params
    optimizer = int(optimizer_multiplier * weights)  # optimizer state (moments, etc.)
    activations = int(bytes_per_activation_per_sample * batch_size)
    scratch = int(scratch_multiplier * (weights + activations))

    total_bytes = weights + gradients + optimizer + activations + scratch

    def to_gb(b): return b / (1024 ** 3)

    breakdown_bytes = {
        "weights": weights,
        "gradients": gradients,
        "optimizer": optimizer,
        "activations": activations,
        "scratch": scratch,
    }
    breakdown_gb = {k: to_gb(v) for k, v in breakdown_bytes.items()}

    return {
        "total_gb": to_gb(total_bytes),
        "breakdown_bytes": breakdown_bytes,
        "breakdown_gb": breakdown_gb,
    }

# Example:
# estimate_gpu_memory_gb(P=1_000_000_000, bytes_per_activation_per_sample=2000, batch_size=32, optimizer_multiplier=2, bytes_per_param=4, scratch_multiplier=0.1)

nvidia-smi --query-gpu=index,name,utilization.gpu,utilization.memory,memory.total,memory.used,temperature.gpu --format=csv -l 2

nsys profile -o trace --trace=nvtx,cuda,osrt --capture-range=nvtx ./train.sh

ncu --set full ./train_kernel

with torch.profiler.profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as p:
    model(input)
p.export_chrome_trace("trace.json")

import triton
@triton.jit
def fused_attn_kernel(Q, Kt, V, Out, stride_q, stride_kt, stride_v, stride_out, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_q = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_k = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    # load tiles
    q = tl.load(Q + offs_q[:, None]*stride_q + tl.arange(0,BLOCK_K)[None,:]*...)
    kt = tl.load(Kt + ... )
    # compute QK^T with triton.dot (uses Tensor Cores when possible)
    logits = tl.dot(q, kt)
    logits = logits / tl.sqrt(BLOCK_K)
    m = tl.max(logits, 1)
    exp = tl.exp(logits - m[:,None])
    denom = tl.sum(exp, 1)
    p = exp / denom[:,None]
    v = tl.load(V + ...)
    out_tile = tl.dot(p, v)
    tl.store(Out + ..., out_tile)