code/gpu/mem_probe.py

"""Single-variant CUDA memory probe (Claims 2 & 3). ONE variant per process,
per the authors' own benchmarking guidance (flashoptim issue #1).

Prints exactly one JSON line: bytes for params / grads / optimizer state and
steady-state peak, plus bytes-per-parameter for each.

Variants:
  ref        fp32 params + torch.optim.AdamW + bf16 autocast fwd/bwd (paper Reference)
  flash      cast_model bf16 + FlashAdamW(master_weight_bits=24, quantize=True)
  flash_gr   flash + enable_gradient_release
  split_only cast_model bf16 + FlashAdamW(quantize=False, master_weight_bits=24)
  quant_only fp32 params + FlashAdamW(quantize=True)  (master bits N/A on fp32)
"""

import argparse
import gc
import json
import sys

import torch


def build_model(name: str):
    if name == "resnet50":
        import torchvision
        m = torchvision.models.resnet50()
        x = torch.randint(0, 255, (8, 3, 224, 224), dtype=torch.uint8)
        inp = (x.cuda().to(torch.float32) / 255.0,)
        tgt = torch.randint(0, 1000, (8,)).cuda()
        loss_fn = lambda out: torch.nn.functional.cross_entropy(out, tgt)
        fwd = lambda model: loss_fn(model(inp[0].to(next(model.parameters()).dtype)))
        return m, fwd
    from models import GPT, SIZES
    c = SIZES[name]
    m = GPT(c)
    ix = torch.randint(0, c.vocab_size, (4, 512)).cuda()
    tg = torch.randint(0, c.vocab_size, (4, 512)).cuda()

    def fwd(model):
        _, loss = model(ix, tg)
        return loss
    return m, fwd


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--variant", required=True,
                    choices=["ref", "flash", "flash_gr", "split_only", "quant_only"])
    ap.add_argument("--model", required=True)
    args = ap.parse_args()

    torch.manual_seed(0)
    dev = torch.device("cuda")
    model, fwd = build_model(args.model)
    n_params = sum(p.numel() for p in model.parameters())

    handle = None
    if args.variant in ("flash", "flash_gr", "split_only"):
        from flashoptim import FlashAdamW, cast_model, enable_gradient_release
        model = model.to(dev)
        cast_model(model, dtype=torch.bfloat16)
    else:
        model = model.to(dev)  # fp32
    gc.collect(); torch.cuda.empty_cache(); torch.cuda.synchronize()
    bytes_params = torch.cuda.memory_allocated()

    if args.variant == "ref":
        opt = torch.optim.AdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95),
                                weight_decay=0.1)
    elif args.variant in ("flash", "flash_gr"):
        opt = FlashAdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95),
                         weight_decay=0.1, master_weight_bits=24, quantize=True)
    elif args.variant == "split_only":
        opt = FlashAdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95),
                         weight_decay=0.1, master_weight_bits=24, quantize=False)
    elif args.variant == "quant_only":
        from flashoptim import FlashAdamW  # fp32 params
        opt = FlashAdamW(model.parameters(), lr=1e-4, betas=(0.9, 0.95),
                         weight_decay=0.1, master_weight_bits=None, quantize=True)
    if args.variant == "flash_gr":
        handle = enable_gradient_release(model, opt)

    use_autocast = args.variant in ("ref", "quant_only")

    def one_iter():
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast):
            loss = fwd(model)
        loss.backward()
        opt.step()
        opt.zero_grad(set_to_none=True)
        return float(loss.detach())

    # warmup: materialize grads + optimizer states + JIT kernels (3 iters)
    losses = [one_iter() for _ in range(3)]
    torch.cuda.synchronize()

    # grads freed by zero_grad(set_to_none) -> allocated == params + opt states
    gc.collect(); torch.cuda.empty_cache(); torch.cuda.synchronize()
    bytes_params_plus_state = torch.cuda.memory_allocated()
    bytes_state = bytes_params_plus_state - bytes_params

    # grads: measure right after a fresh backward, before step
    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_autocast):
        loss = fwd(model)
    loss.backward()
    del loss
    gc.collect(); torch.cuda.synchronize()
    if args.variant == "flash_gr":
        bytes_grads = 0  # released eagerly during backward
    else:
        bytes_grads = torch.cuda.memory_allocated() - bytes_params_plus_state
    opt.step(); opt.zero_grad(set_to_none=True)

    # steady-state peak over one full iteration
    torch.cuda.reset_peak_memory_stats()
    losses.append(one_iter())
    torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated()

    if handle is not None:
        handle.remove()
    print(json.dumps(dict(
        variant=args.variant, model=args.model, n_params=n_params,
        bytes_params=bytes_params, bytes_state=bytes_state,
        bytes_grads=bytes_grads, peak_steady=peak,
        bpp_params=bytes_params / n_params, bpp_state=bytes_state / n_params,
        bpp_grads=bytes_grads / n_params,
        bpp_params_state_grads=(bytes_params + bytes_state + bytes_grads) / n_params,
        losses=[round(x, 4) for x in losses])))
    sys.stdout.flush()


if __name__ == "__main__":
    main()