Lesson 70: PyTorch Hooks, Profiling & torch.compile

Discussion prompt

Before we open Lesson 70: PyTorch Hooks, Profiling & torch.compile: without looking back, what was the main idea of kNN and the Curse of Dimensionality, and what could you do by the end of it that you could not do before?

Hint: One sentence for the idea, one for the skill. If the second one is blank, that is the part to revisit.

Comparison matrix

From What a forward hook is: refill the signature column from what you know. The rest of the table is as it appeared.

Predict first

Build a tiny CNN (conv1 -> ReLU -> conv2 -> pool -> fc) and register forward hooks on both conv layers.

Commit before you compute: what does Hooking a CNN to capture activations come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Call h.remove() after each forward pass you need — or use the context manager form

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Uncleaned hooks accumulate: the same hook fires again every forward pass and you collect stale references.

Comparison matrix

From Hooking a CNN to capture activations: every row here is a choice with a cost. Fill the activation shape column, then say which row you would actually pick and what you give up for it.

Predict first

A student writes this, and it looks reasonable:

Register register_forward_hook once inside a training loop body to log activations each step.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: Each iteration registers a NEW hook without removing the previous one.

Register once before the loop; remove immediately after sampling.

Why: Each iteration registers a NEW hook without removing the previous one. After 100 epochs you have 100 hooks all firing, storing 100 copies of the activation tensor.

Discussion prompt

register_full_backward_hook fires during the backward pass with (module, grad_input, grad_output) — both are tuples of tensors (or None if detached).

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Discussion prompt

Attach full-backward hooks to conv1, conv2, and fc. Run one forward + backward and read off the gradient norms flowing into each layer.

What do you need to know — or decide — before the first line can be written? List everything the problem has to hand you.

Hint: Anything you would have to invent to get started is a thing the problem must supply.

Discussion prompt

Work backwards. The example finished here:

Gradient norm drops from 3.16 at fc to 0.047 at conv1 — a 67x reduction across only 3 layers

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Discussion prompt

Explain Three memory levers by analogy to something with no Machine Learning in it at all — a queue, a recipe, a map, a bank balance, whatever fits. Then say where your analogy breaks.

Hint: An analogy that never breaks is not an analogy, it is the same idea wearing a hat. Find the seam — that is the part that is actually new.

Predict first

An 8-block network, batch=32, hidden=64: each activation tensor is 32x64x4 bytes = 8 KB. Compare storing all vs checkpointing at sqrt(8)=2 segments.

Commit before you compute: what does Gradient checkpointing: verified memory math come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Gradient norms are identical with and without checkpointing

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. Checkpointing only changes WHEN the activations are computed (re-computed on demand during backward), not what gradients are produced.

Discussion prompt

Work backwards. The example finished here:

Gradient norms are identical with and without checkpointing

What was it asked to do, and what must it have been given? Reconstruct the problem from its answer.

Hint: Every quantity in the result had to enter somewhere. Account for each one.

Predict first

A student writes this, and it looks reasonable:

The GPU runs out of memory mid-training. Call torch.cuda.empty_cache() to free it.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: empty_cache() only releases the caching allocator's free blocks — blocks PyTorch already freed but hasn't returned to the OS yet.

First del the tensors you no longer need, then call empty_cache() if needed.

Why: empty_cache() only releases the caching allocator's free blocks — blocks PyTorch already freed but hasn't returned to the OS yet. It cannot free tensors that are still referenced. If you're OOM, the problem is live tensors, not the cache.

Discussion prompt

The rule this trap just fixed:

del removes the Python reference; if refcount hits 0, PyTorch's allocator marks the block free. Only then does empty_cache() have something to return to the OS.

Now break it on purpose. Build a case that violates it and follow the consequences until something visibly fails. Where does the failure first show up — and would you have noticed it if you had not been looking?

Hint: The dangerous rules are the ones whose violation still produces an answer. If yours fails loudly, try to find one that fails quietly.

Discussion prompt

Explain torch.profiler: find the bottleneck to a student a year behind you. No notation, no jargon they have not met — and it still has to be true.

Hint: If your explanation needs a symbol they have never seen, you are describing the notation rather than the idea.

Predict first

Profile 5 training steps on a 2-layer MLP (128 -> 256 -> 10). Sort by self CPU time to find the top op.

Commit before you compute: what does Profiling a training loop come out to? A rough magnitude and the right form is enough — the point is to have something concrete to be wrong about.

Correct: Loss backward (nll_loss_backward + _log_softmax_backward) takes as much time as the forward pass

Why: A prediction you can defend turns the computation into a check rather than a leap of faith — and an answer that contradicts it is caught on the spot. On small models, backward is ~2x the forward cost — consistent with the rule of thumb.

Sort into buckets

Sort these by calls, from memory, without looking back at Profiling a training loop. Telling them apart on the spot is the skill; the table is only where the answer happens to be written down.

5

aten::nll_loss_backward; aten::_log_softmax

10

aten::addmm (linear fwd)

15

aten::mm (grad matmul)

g1

calls is "5" for aten::nll_loss_backward, aten::_log_softmax — that is what the table on "Profiling a training loop" records, and it is the single property separating this group from the rest.

g2

calls is "10" for aten::addmm (linear fwd) — that is what the table on "Profiling a training loop" records, and it is the single property separating this group from the rest.

g3

calls is "15" for aten::mm (grad matmul) — that is what the table on "Profiling a training loop" records, and it is the single property separating this group from the rest.

Comparison matrix

From torch.compile (PyTorch 2.0+): refill the compile cost column from what you know. The rest of the table is as it appeared.

Discussion prompt

Explain What compile does internally by analogy to something with no Machine Learning in it at all — a queue, a recipe, a map, a bank balance, whatever fits. Then say where your analogy breaks.

Hint: An analogy that never breaks is not an analogy, it is the same idea wearing a hat. Find the seam — that is the part that is actually new.

Predict first

A student writes this, and it looks reasonable:

To compile model inference, call torch.compile on the model object each iteration.

It is wrong. Say what breaks — and say it before you turn the page.

Correct: torch.compile is not a decorator that caches the compiled model — it returns a new compiled wrapper object.

Compile once before the loop; reuse the compiled wrapper.

Why: torch.compile is not a decorator that caches the compiled model — it returns a new compiled wrapper object. Calling it in a loop triggers recompilation every iteration, multiplying compile latency by the number of steps.

Sort into buckets

Some of these hold up and some are the exact mistakes this lesson is built to prevent. Sort them.

Holds up

register_full_backward_hook fires during the backward pass with (module, grad_input, grad_output) — both are tuples of tensors (or None if detached).; On CPU (Windows): inductor requires a C++ compiler (cl.exe). On Linux/CUDA the stack is fully supported out of the box — the typical production target.; Build rules: remove every hook via .remove() after use; don't register hooks inside the training loop; compare grad norms with and without checkpointing.

Breaks

Register register_forward_hook once inside a training loop body to log activations each step.; The GPU runs out of memory mid-training. Call torch.cuda.empty_cache() to free it.

sound

These are stated as this lesson states them — each one survives the edge cases Lesson 70: PyTorch Hooks, Profiling & torch.compile puts it through.

flawed

Each of these is lifted from a trap in this deck: reasonable-sounding, and wrong in a way that only shows up once you rely on it.

Put in order

These are the steps of The hooks + profiling + compile recipe, scrambled. Put them back in order before the next slide shows you.

1. Forward hook: h = layer.register_forward_hook(fn); always call h.remove() when done
2. Backward hook: register_full_backward_hook (not the deprecated register_backward_hook)
3. Memory: del tensor before torch.cuda.empty_cache(); use gradient checkpointing for deep nets
4. Profile: wrap training steps in torch.profiler.profile(...); sort by self_cpu_time_total
5. Compile: compiled = torch.compile(model, mode='default') once before the loop; first call pays compile cost

Why: This is the order the recipe itself gives. Recalling the sequence without the slide in front of you is the difference between recognising the method and being able to run it — most of what goes wrong in practice is a step done out of turn.

Discussion prompt

The hooks + profiling + compile recipe works on the cases you have just seen. Push it to the edge: what is the most degenerate input it still handles — empty, zero, one item, everything equal — and what is the first case where it stops being true? Name the case, not just "it breaks".

Hint: Try the smallest legal input, then the largest, then the one where two things collide. Methods are specified at their edges; the middle takes care of itself.

Eliminate the wrong options

A register_forward_hook callback receives which arguments?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. (module, input, output)
• B. (module, grad_input, grad_output)
• C. (loss, output, target)
• D. (input, output)

Survives elimination: A

Why: Forward hooks receive (module, input, output): the module that fired, its input tuple, and its output tensor. The grad_input/grad_output signature belongs to backward hooks.

Check your understanding

A register_forward_hook callback receives which arguments?

• A. (module, input, output) (correct)
• B. (module, grad_input, grad_output)
• C. (loss, output, target)
• D. (input, output)

Answer: A

Why: Forward hooks receive (module, input, output): the module that fired, its input tuple, and its output tensor. The grad_input/grad_output signature belongs to backward hooks.

Why B tempts people

That is the signature of register_full_backward_hook, not forward. Forward hooks fire after forward(); backward hooks fire after backward().

Why C tempts people

No PyTorch hook receives the loss or target directly; hooks are attached to layers, not to loss functions.

Why D tempts people

The module itself is always the first argument — hooks need it to identify which layer fired, especially when the same function is registered on multiple layers.

Predict first

Gradient checkpointing reduces peak activation memory by not storing intermediate tensors. How does backward still get the activations it needs for gradient computation?

Answer it in your own words, now, with nothing to choose from. The options are on the next slide — and picking the right one off a list is an easier skill than producing it.

Correct: It re-runs the forward pass over each checkpointed segment during backward

Why: The backward pass re-runs the forward computation for each checkpoint segment to regenerate the activations it needs. This trades ~2x forward compute for O(sqrt(L)) instead of O(L) activation memory.

Check your understanding

Gradient checkpointing reduces peak activation memory by not storing intermediate tensors. How does backward still get the activations it needs for gradient computation?

• A. It re-runs the forward pass over each checkpointed segment during backward (correct)
• B. It stores all activations on CPU instead of GPU
• C. It approximates the activations using the input only
• D. It skips gradient computation for checkpointed layers

Answer: A

Why: The backward pass re-runs the forward computation for each checkpoint segment to regenerate the activations it needs. This trades ~2x forward compute for O(sqrt(L)) instead of O(L) activation memory.

Why B tempts people

CPU offloading is a separate technique (e.g., Fairscale's offload). Vanilla torch.utils.checkpoint keeps everything in the same memory tier and simply does not store tensors.

Why C tempts people

Approximations would produce wrong gradients — PyTorch never approximates. The recomputed values are exact.

Why D tempts people

Skipping gradients would break backprop entirely. All layers still receive correct gradients; only the timing of activation computation changes.

Eliminate the wrong options

You are deploying a small model called thousands of times per second. Which torch.compile mode most directly addresses the per-call Python/CUDA launch overhead?

3 of these 4 are wrong. Strike them one at a time, and say what rules each one out before you strike the next. The survivor is the answer.

• A. reduce-overhead
• B. max-autotune
• C. default
• D. eager (no compile)

Survives elimination: A

Why: reduce-overhead explicitly targets the per-call Python and CUDA dispatch cost, which dominates when the model is small and called very frequently. max-autotune optimizes kernel selection but carries a high compilation cost that is only worth paying for large compute-bound models.

Check your understanding

You are deploying a small model called thousands of times per second. Which torch.compile mode most directly addresses the per-call Python/CUDA launch overhead?

• A. reduce-overhead (correct)
• B. max-autotune
• C. default
• D. eager (no compile)

Answer: A

Why: reduce-overhead explicitly targets the per-call Python and CUDA dispatch cost, which dominates when the model is small and called very frequently. max-autotune optimizes kernel selection but carries a high compilation cost that is only worth paying for large compute-bound models.

Why B tempts people

max-autotune searches exhaustively for the fastest kernel at high compile cost; it is best for large models where kernel selection matters, not for small models where overhead is the bottleneck.

Why C tempts people

default is a good general choice but is not specifically tuned for the launch-overhead use case that high-QPS small-model serving creates.

Why D tempts people

Without compile, every call pays full Python dispatcher overhead. For thousands of calls per second, that overhead adds up to significant latency.

Discussion prompt

Build rules: remove every hook via .remove() after use; don't register hooks inside the training loop; compare grad norms with and without checkpointing.

That is stated as though it always holds. Do one of two things: produce a case where it fails, or say precisely what rules such a case out. "It just does" is not on the menu.

Hint: Hunt at the extremes first — zero, one, negative, empty, equal. If every extreme survives, the reason they survive is the proof.

Comparison matrix

From Milestone 2 — backward hook gradient flow: every row here is a choice with a cost. Fill the grad norm column, then say which row you would actually pick and what you give up for it.

Comparison matrix

From Milestone 3 — profile + checkpoint: refill the self CPU (5 steps) column from what you know. The rest of the table is as it appeared.

Draw it

One page, no notation unless you need it: draw how these connect — Forward hooks: intercept activations · Backward hooks: gradient flow · Memory management & checkpointing · Profiling & torch.compile · Your turn: hook it, profile it, compile it. Put an arrow wherever one of them is what makes another possible, and label the arrow with why.