24. Profiling a Forward Pass for a Memory Bottleneck

Problem

Two implementations of the "same" block can have identical outputs and wildly different activation memory footprints -- which matters a lot when deciding which one to actually ship. Rather than guessing, the standard technique is to register a forward hook on every leaf submodule, run one real forward pass, and record how many bytes each submodule's output tensor occupies; whichever submodule produced the single largest intermediate tensor is that implementation's memory bottleneck.

Implement a profiler that takes two candidate nn.Module implementations and a shared input, runs a forward pass through each with forward hooks attached to every leaf submodule, and reports each implementation's peak single-layer activation footprint in bytes plus which of the two has the larger one.

Source: src/24_profile_forward_bottleneck.py

def profile_forward_bottleneck(
    impl_a: torch.nn.Module, impl_b: torch.nn.Module, x: torch.Tensor,
) -> dict: ...

>>> import torch.nn as nn
>>> wide = nn.Sequential(nn.Linear(8, 64), nn.ReLU(), nn.Linear(64, 8))
>>> narrow = nn.Sequential(nn.Linear(8, 16), nn.ReLU(), nn.Linear(16, 8))
>>> report = profile_forward_bottleneck(wide, narrow, torch.randn(4, 8))
>>> report["bottleneck"]
'a'

Step-by-Step Approach

1. Write a helper that profiles one model: walk model.named_modules() and keep only the leaves (modules with no children) -- container modules like nn.Sequential itself would double-count their children's output.
2. Register a forward hook on each leaf that records output.numel() * output.element_size() into a {name: bytes} dict, keyed by that hook's captured module name.
3. Run one forward pass, wrapped in torch.no_grad() since this is pure profiling, not training.
4. Remove every hook in a finally block, so a raised exception during the forward pass can't leave stale hooks attached for the next profiling call on the same model.
5. Find the entry with the largest byte count using max(..., key=...); because dicts preserve insertion order and max keeps the first maximal item under a strict "greater than" comparison, this naturally implements the "first submodule wins ties" rule.
6. Run the helper on both implementations, then compare the two peak byte counts to decide which one is the bottleneck.

The key insight is that activation memory is dominated by whichever single intermediate tensor is largest, not by the sum of all intermediates -- PyTorch frees most intermediates as soon as backward no longer needs them, so profiling per-layer peak size (not total) is what actually predicts which implementation will run out of memory first.

Reference solution

def _peak_activation(model: nn.Module, x: torch.Tensor) -> tuple[str, int]:
    sizes: dict[str, int] = {}

    def make_hook(name):
        def hook(_module, _inputs, output):
            # numel() * element_size() = exact byte footprint of this one
            # activation tensor, independent of dtype (fp32/fp16/bf16 alike).
            sizes[name] = output.numel() * output.element_size()
        return hook

    handles = [
        module.register_forward_hook(make_hook(name))
        for name, module in model.named_modules()
        if len(list(module.children())) == 0  # leaf modules only
    ]

    try:
        with torch.no_grad():
            model(x)
    finally:
        # hooks must come off even if the forward pass raises
        for handle in handles:
            handle.remove()

    # keeps the FIRST maximal entry -- matches the tie-break rule
    peak_name, peak_bytes = max(sizes.items(), key=lambda kv: kv[1])
    return peak_name, peak_bytes


def profile_forward_bottleneck(
    impl_a: nn.Module, impl_b: nn.Module, x: torch.Tensor,
) -> dict:
    a_layer, a_bytes = _peak_activation(impl_a, x)
    b_layer, b_bytes = _peak_activation(impl_b, x)
    return {
        "a_peak_bytes": a_bytes, "a_peak_layer": a_layer,
        "b_peak_bytes": b_bytes, "b_peak_layer": b_layer,
        "bottleneck": "b" if b_bytes > a_bytes else "a",
    }

Key Functions & Tricks

• module.register_forward_hook(fn) — attaches fn(module, inputs, output) to run every time that module's forward completes, without modifying the module's code.
• model.named_modules() — recursively yields every submodule with its dotted name; filtering to leaves avoids double-counting container modules.
• output.numel() * output.element_size() — the exact byte size of a tensor, correct for any dtype (a common bug is hardcoding * 4 for float32 and silently breaking under mixed precision).
• handle.remove() — every hook returned by register_forward_hook must be explicitly removed, or it keeps firing (and accumulating) on every future forward pass through that model.
• try/finally around the forward pass — guarantees hook cleanup even when profiling a model that might raise.

How to Recognize This Pattern