The instruments that answer nothing here

Run every CUDA performance call on this laptop and the failures sort into three kinds. torch.cuda.memory_allocated() returns 0. torch.cuda.synchronize() raises. torch.cuda.amp.GradScaler() warns once and disables itself, then behaves like an object that works.

The first kind is the dangerous one. A memory report that prints 0 bytes allocated and 0 reserved looks exactly like a program with nothing on the device, and memory_stats() returning an empty dictionary is what makes memory_summary() raise a KeyError a line later rather than say anything useful. Nothing in that sequence tells you the answer was never about your program.

So gate on torch.cuda.is_available() yourself, in the reporting code, and print the gate's result next to the numbers. A performance report from a CPU-only box that quietly contains a CUDA memory section is not wrong in its arithmetic. It is answering a question nobody asked.

An instrument that cannot see your device will still print a number.

torch 2.2.2 | cuda available: False
memory_allocated()                   0
max_memory_allocated()               0
memory_reserved()                    0
memory_stats()                       OrderedDict()
memory_summary()                     KeyError: 'allocated_bytes.all.current'
synchronize()                        AssertionError: Torch not compiled with CUDA enabled
Stream()                             RuntimeError: Tried to instantiate dummy base class Stream
Event(enable_timing=True)            RuntimeError: Tried to instantiate dummy base class Event
GradScaler().is_enabled()            False  [warned: torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.]
autocast(cuda) on a matmul           torch.float32  [warned: User provided device_type of 'cuda', but CUDA is not available. Disabling]
profile(CPU and CUDA)                'ran, no CUDA columns'  [warned: CUDA is not available, disabling CUDA profiling]

import warnings

import torch
from torch.profiler import ProfilerActivity, profile

def probe(label, f):
    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        try:
            out = repr(f())
        except Exception as err:
            out = f"{type(err).__name__}: {err}"
    note = f"  [warned: {w[0].message}]" if w else ""
    print(f"{label:36s} {out}{note}")

def cuda_activity():
    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as p:
        (torch.randn(64, 64) @ torch.randn(64, 64)).sum()
    return "CUDA columns" if "Self CUDA" in p.key_averages().table(row_limit=1) else "ran, no CUDA columns"

def autocast_matmul():
    with torch.autocast("cuda", dtype=torch.float16):
        return (torch.ones(4, 4) @ torch.ones(4, 4)).dtype

print("torch", torch.__version__, "| cuda available:", torch.cuda.is_available())
probe("memory_allocated()", torch.cuda.memory_allocated)
probe("max_memory_allocated()", torch.cuda.max_memory_allocated)
probe("memory_reserved()", torch.cuda.memory_reserved)
probe("memory_stats()", torch.cuda.memory_stats)
probe("memory_summary()", lambda: torch.cuda.memory_summary()[:40])
probe("synchronize()", torch.cuda.synchronize)
probe("Stream()", torch.cuda.Stream)
probe("Event(enable_timing=True)", lambda: torch.cuda.Event(enable_timing=True))
probe("GradScaler().is_enabled()", lambda: torch.cuda.amp.GradScaler().is_enabled())
probe("autocast(cuda) on a matmul", autocast_matmul)
probe("profile(CPU and CUDA)", cuda_activity)

A CUDA stream is a queue with one rule: operations inside it run in the order they were issued. Every device has a default stream, and if you never make another one, everything you have ever written is already ordered against itself, which is why the async dispatch the chapter describes is safe by default rather than a race.

Two streams have no such rule between them. The docs' own broken example queues a normal_() on the default stream and a sum on a side stream, and the sum may start before the fill finishes. Fixing it takes two calls: wait_stream makes the side stream wait for what the default stream already queued, and record_stream keeps the source tensor alive until the side stream is done with it. The second one is easy to forget because the failure it prevents is the caching allocator handing your buffer to somebody else while a kernel is still reading it.

For timing, the same asynchrony means the host clock is measuring the wrong thing, and CUDA events are the device-side answer. Record an event before and after, synchronize once, then ask for elapsed_time between them, and the interval you get was measured by the device rather than by Python. Lesson one's Timer does the coarser version of this for you by synchronizing before each read.

None of this ran here. Both torch.cuda.Stream() and torch.cuda.Event() raise on this machine, so the block below is the pinned source rather than a capture. No lab on this site measures stream overlap either, because the torch labs run on a plain CPU or on a TPU and a stream is a CUDA object. That leaves the measurement to you, on the first CUDA device you get hold of; the exercises below say what to record.

# lines 229-237: timing with events instead of the host clock
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()

# Run some things here

end_event.record()
torch.cuda.synchronize()  # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)

# lines 257-262: the docs call this one incorrect
cuda = torch.device("cuda")
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
with torch.cuda.stream(s):
    # sum() may start execution before normal_() finishes!
    B = torch.sum(A)

# lines 270-275: the fixed version
cuda = torch.device("cuda")
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
s.wait_stream(torch.cuda.default_stream(cuda))  # NEW!
with torch.cuda.stream(s):
    B = torch.sum(A)
A.record_stream(s)  # NEW!

The reason to own a second stream at all is usually a copy. A host-to-device transfer can run while the device computes, but only if three things are true at once: the host buffer is page-locked, the copy is issued with non_blocking=True, and the copy is not queued behind the compute it was supposed to overlap with.

Page-locking is the loader's job and it has its own home. What pin_memory=True does, what it silently does not do on a machine with no accelerator, and the thread it starts when it does apply are told at /pytorch/data/same-order-different-numbers. The half that belongs here is that pinning alone buys nothing: an unpinned copy with non_blocking=True is a synchronous copy wearing an argument.

You can watch the argument mean nothing right now. A host-to-host copy accepts non_blocking=True, returns an unpinned tensor, and warns about none of it, because there is no device for the transfer to be asynchronous with respect to. Every flag in this section behaves that way here, accepted and inert, which is what the first section of this lesson sorted into three kinds.

One asymmetry is worth carrying forward for when you do have a device. Backward ops run on the stream their forward op ran on, so a forward that splits work across streams gets a backward that splits the same way without you arranging it, and gradients produced inside a stream context must be consumed inside it or after an explicit wait.

import torch

x = torch.ones(1024)
y = x.to("cpu", non_blocking=True)
print("cpu to cpu, non_blocking:", tuple(y.shape), y.dtype, "| is_pinned:", y.is_pinned())

# ---- stdout ----
# cpu to cpu, non_blocking: (1024,) torch.float32 | is_pinned: False

# and the rule it is half of, verbatim from docs/source/notes/cuda.rst at v2.2.2:
#
#   Host to GPU copies are much faster when they originate from pinned
#   (page-locked) memory. [...] Once you pin a tensor or storage, you can use
#   asynchronous GPU copies. Just pass an additional non_blocking=True argument
#   to a to() or a cuda() call.

Under torch.autocast, an op is not asked what dtype it wants. It is looked up in two lists. One list casts its inputs down to the low-precision dtype, one casts them up to float32, and anything in neither list runs in whatever dtype its inputs already have. That third case is most ops, and it is why a bf16 value flows onward through a chain of unlisted ops without anyone deciding it should.

Run the lookup on this machine with torch.autocast("cpu", dtype=torch.bfloat16) and the shape is visible in one table. Matmul, linear and conv come back bfloat16. mse_loss and prod come back float32. relu comes back float32 when handed a float32 input and bfloat16 when handed the output of a matmul, because relu is in neither list.

The CPU lists and the CUDA lists are different documents, and that is where an audit run here stops transferring. log_softmax, softmax, sum and layer_norm are all in the CUDA float32 list and in neither CPU list, so the same script under autocast produces float32 for those on a GPU and bfloat16 here. A dtype audit run on CPU is a real audit of the CPU table and tells you nothing certain about the other one.

The weights never change. lin.weight.dtype is still float32 during and after the region, because autocast casts at the operator boundary rather than converting your parameters, and it caches that cast so a weight used repeatedly inside one region is converted once. Profile five calls of one Linear under autocast and there are seven aten::_to_copy calls: five inputs, plus the weight and bias converted once between them.

import torch
import torch.nn as nn
from torch.profiler import ProfilerActivity, profile

lin = nn.Linear(512, 512)
x = torch.randn(64, 512)

with torch.autocast("cpu", dtype=torch.bfloat16):
    lin(x)                       # warm the cast cache outside the capture

with profile(activities=[ProfilerActivity.CPU]) as prof:
    with torch.autocast("cpu", dtype=torch.bfloat16):
        for _ in range(5):
            lin(x)

counts = {e.key: e.count for e in prof.key_averages()}
print("aten::addmm     ", counts["aten::addmm"])
print("aten::_to_copy  ", counts["aten::_to_copy"])
print("weight dtype    ", lin.weight.dtype)

# ---- stdout ----
# aten::addmm      5
# aten::_to_copy   7
# weight dtype     torch.float32

import torch
import torch.nn.functional as F

x = torch.randn(8, 64)
w = torch.randn(64, 64)
img = torch.randn(2, 3, 16, 16)
k = torch.randn(4, 3, 3, 3)

cases = {
    "x @ w": lambda: x @ w,
    "F.conv2d": lambda: F.conv2d(img, k),
    "F.linear": lambda: F.linear(x, w),
    "relu, float32 in": lambda: torch.relu(x),
    "relu, bfloat16 in": lambda: torch.relu(x @ w),
    "sum": lambda: (x @ w).sum(),
    "softmax": lambda: torch.softmax(x @ w, -1),
    "log_softmax": lambda: torch.log_softmax(x @ w, -1),
    "layer_norm": lambda: F.layer_norm(x @ w, (64,)),
    "mse_loss": lambda: F.mse_loss(x @ w, torch.randn(8, 64)),
    "prod": lambda: (x @ w).prod(),
}
with torch.autocast("cpu", dtype=torch.bfloat16):
    for name, f in cases.items():
        print(f"{name:20s} {f().dtype}")

# ---- stdout ----
# x @ w                torch.bfloat16
# F.conv2d             torch.bfloat16
# F.linear             torch.bfloat16
# relu, float32 in     torch.float32
# relu, bfloat16 in    torch.bfloat16
# sum                  torch.bfloat16
# softmax              torch.bfloat16
# log_softmax          torch.bfloat16
# layer_norm           torch.bfloat16
# mse_loss             torch.float32
# prod                 torch.float32

Autocast is a dtype policy, not a speedup. Whether the policy pays depends entirely on whether the hardware has arithmetic for the dtype it picked, and this machine is a clean demonstration of the failure case: an Intel CPU whose top vector capability is AVX2, with no bfloat16 instructions underneath.

The same Linear that runs in 0.192 milliseconds in float32 takes 8.198 under autocast("cpu", dtype=torch.bfloat16). Casting the module and the input to bfloat16 by hand, with no autocast anywhere, takes 7.783, so the casts are a small part of it and the arithmetic path is the rest. Between thirty and forty times slower, on a change that a GPU or TPU benchmark would report as a speedup.

Nothing about that contradicts the chapter. It is the same claim from the other side: the dtype is fast where the tensor cores or MXUs implement it, and this box implements none of them. State the chip whenever you quote a dtype result, because a bf16 number without a chip beside it is not interpretable at all.

The scaler behaves the same way. GradScaler exists for float16's narrow exponent range, and bfloat16 does not need it. On this machine it also does not run: constructing one warns and disables it, get_scale() reads 1.0, and scaler.scale(loss).backward() produces exactly the gradients an unscaled backward would. Code written that way is not wrong here, it is inert here, and a CPU run proves nothing about whether your scaler is doing its job.

import warnings

import torch

with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    scaler = torch.cuda.amp.GradScaler()
    print("warnings:", [str(x.message) for x in w])
print("enabled:", scaler.is_enabled(), "| scale:", scaler.get_scale())

lin = torch.nn.Linear(4, 4)
opt = torch.optim.SGD(lin.parameters(), lr=0.1)
loss = lin(torch.randn(2, 4)).sum()
scaler.scale(loss).backward()
print("bias grad:", lin.bias.grad)
scaler.step(opt)
scaler.update()
print("new scale:", scaler.get_scale())

# ---- stdout ----
# warnings: [torch.cuda.amp.GradScaler is enabled, but CUDA is not available.  Disabling.]
# enabled: False | scale: 1.0
# bias grad: tensor([2., 2., 2., 2.])
# new scale: 1.0

PyTorch does not hand memory back to CUDA when your tensor dies. It keeps it, in a caching allocator, so the next allocation of that size costs nothing and no device synchronization is needed to free anything. Two numbers follow from that, and confusing them is the most common way a memory report becomes fiction.

memory_allocated is what your live tensors hold. memory_reserved is what the allocator has taken from the driver, cached blocks included, and it is the number nvidia-smi shows. A run whose reserved figure is far above its allocated figure is not leaking; it is holding a cache. empty_cache() returns the unused part of that cache to the driver and changes nothing about what your tensors hold.

memory_stats() is where the diagnosis lives, and four of its keys are worth knowing before you need them. requested_bytes against allocated_bytes shows what allocation rounding costs you. num_alloc_retries counts the times a cudaMalloc failed and forced a cache flush, which is fragmentation showing itself before the crash does. num_ooms counts the crashes. Each core statistic carries current, peak, allocated and freed, so the peak is available without you polling for it.

For the question of which line allocated what, the snapshot tools are the instrument: _record_memory_history turns on stack capture per allocation, _dump_snapshot writes a pickle, and the viewer at pytorch.org/memory_viz renders it locally. On this machine every one of those calls answers about a device that does not exist, so that workflow waits for a CUDA device too. The CPU substitute you do have is the profiler's memory columns from lesson two, which answer the same shape of question about the host allocator.

- ``"allocated.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
  number of allocation requests received by the memory allocator.
- ``"allocated_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
  amount of allocated memory.
- ``"segment.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
  number of reserved segments from ``cudaMalloc()``.
- ``"reserved_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
  amount of reserved memory.

- ``"num_alloc_retries"``: number of failed ``cudaMalloc`` calls that
  result in a cache flush and retry.
- ``"num_ooms"``: number of out-of-memory errors thrown.

- ``"requested_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``:
  memory requested by client code, compare this with allocated_bytes to check if
  allocation rounding adds too much overhead.