Add capture_time_hooks to make_graphed_callables for non-capturable per-callable hooks

[buptzyb]

Description

Summary

make_graphed_callables previously had no mechanism to run per-callable
hooks during warmup/capture that must execute outside the CUDA graph
capture context (i.e., hooks that are inherently non-capturable, such as
CPU-side state updates or FSDP parameter un-shard/re-shard calls).

This PR adds a new capture_time_hooks parameter to
make_graphed_callables that accepts per-callable hooks invoked at
capture time (warmup iterations and graph capture), but intentionally
executed outside the CUDA graph capture context so they are not
recorded into the graph and will not be replayed.

Changes

• transformer_engine/pytorch/graph.py:
  • Add capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]]
    parameter to make_graphed_callables
  • Invoke hooks around forward and backward passes during both warmup
    iterations and graph capture, in both the _order is not None and
    _order is None capture paths
  • Hook dict structure mirrors PyTorch's _forward_pre_hooks /
    _forward_hooks format: Dict[hook_type, Dict[handle_id, hook_fn]]
    where hook types are forward_pre, forward, pre_backward,
    backward
  • Rename parameter from capture_hooks → capture_time_hooks with
    updated docstring clarifying the non-capturable semantics

Motivation

Used by Megatron-LM's FSDP integration: during CUDA Graph capture,
PyTorch's memory allocator is frozen, causing FSDP parameter
un-shard/re-shard (which requires allocation) to fail. By routing FSDP
hooks through capture_time_hooks, they execute outside the capture
context and are manually driven at the right points during
warmup/capture, while the graph itself only records pure GPU compute.

Hook Invocation Order

For each callable at each warmup/capture iteration:

1. forward_pre hooks — before func(*args, **kwargs)
2. (CUDA graph capture context entered)
3. Forward pass
4. (CUDA graph capture context exited)
5. forward hooks — after forward
6. pre_backward hooks — before torch.autograd.backward
7. Backward pass
8. backward hooks — after backward

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adds a capture_time_hooks parameter to make_graphed_callables, letting callers inject per-callable hooks that run outside the CUDA graph capture context (e.g., FSDP un-shard/re-shard). It also refactors the monolithic warmup loop into _run_warmup_forward / _run_warmup_backward helpers and unifies the _order is None / _order is not None warmup paths.

• P1 – hook output inconsistency: _run_warmup_forward passes flattened outputs to forward hooks, while both capture paths pass the raw (unflattened) return value of func(...). Any hook that inspects output tensors will observe different types depending on the phase.
• P1 – pre_warmup_hook/post_warmup_hook regression: These hooks were previously called once per callable; after the refactor they are called once globally, silently breaking callers that relied on per-callable invocation.

Confidence Score: 3/5

Not safe to merge as-is; two P1 behavioral regressions need fixing before the feature is reliable.

The forward-hook output-type mismatch (flattened vs. raw) and the silently changed pre/post_warmup_hook invocation semantics are both present defects in the changed code that can produce wrong behavior for existing and new callers without any runtime error.

transformer_engine/pytorch/graph.py — specifically the _run_warmup_forward forward-hook call and the pre/post_warmup_hook placement

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Adds capture_time_hooks parameter and refactors warmup into helper functions; two behavioral issues: (1) forward hooks get flattened outputs during warmup but raw outputs during capture, breaking hook API contract; (2) pre/post_warmup_hook are now called once globally instead of once per callable, silently changing existing semantics.

Sequence Diagram

sequenceDiagram
    participant Caller
    participant make_graphed_callables
    participant _run_warmup_forward
    participant _run_warmup_backward
    participant CUDAGraph

    Caller->>make_graphed_callables: capture_time_hooks=[{forward_pre, forward, pre_backward, backward}]
    make_graphed_callables->>make_graphed_callables: pre_warmup_hook() [once globally]

    loop warmup_iter in range(num_warmup_iters)
        make_graphed_callables->>_run_warmup_forward: (func_idx, func)
        _run_warmup_forward->>Caller: forward_pre hooks(func, args, kwargs) [outside capture]
        _run_warmup_forward->>_run_warmup_forward: func(*args, **kwargs)
        _run_warmup_forward->>Caller: forward hooks(func, args, FLATTENED outputs) [outside capture]
        make_graphed_callables->>_run_warmup_backward: (func_idx, func, outputs, warmup_iter)
        _run_warmup_backward->>Caller: pre_backward hooks(func) [outside capture]
        _run_warmup_backward->>_run_warmup_backward: torch.autograd.backward(...)
        _run_warmup_backward->>Caller: backward hooks(func) [outside capture]
    end

    make_graphed_callables->>make_graphed_callables: post_warmup_hook() [once globally]

    note over make_graphed_callables,CUDAGraph: Graph Capture Phase
    make_graphed_callables->>Caller: forward_pre hooks(func, args, kwargs) [outside capture]
    make_graphed_callables->>CUDAGraph: begin capture
    CUDAGraph->>CUDAGraph: func(*args, **kwargs) [recorded]
    make_graphed_callables->>CUDAGraph: end capture
    make_graphed_callables->>Caller: forward hooks(func, args, RAW outputs) [outside capture]
    make_graphed_callables->>Caller: pre_backward hooks(func) [outside capture]
    make_graphed_callables->>CUDAGraph: begin capture
    CUDAGraph->>CUDAGraph: autograd.backward(...) [recorded]
    make_graphed_callables->>CUDAGraph: end capture
    make_graphed_callables->>Caller: backward hooks(func) [outside capture]

Loading

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[greptile-apps]

forward hook receives inconsistent output types between warmup and capture

During warmup, _run_warmup_forward calls _tree_flatten on the raw outputs and then passes the already-flattened list to the forward hook:

outputs, _ = _tree_flatten(func(*args, **kwargs))
# ...
hook(func, args, outputs)   # ← outputs is a flat list of tensors

But in both capture paths (_order is None and _order is not None, lines ~690 and ~913), the forward hook is called with the raw, unflattened return value of func(...):

with _graph_context_wrapper(fwd_graph, pool=mempool):
    outputs = func(*args, **kwargs)   # raw module output
# ...
hook(func, args, outputs)   # ← outputs is the original nested structure

Any hook that inspects or stores output will see different types depending on whether it is called during warmup or during capture. For FSDP re-shard hooks that key off tensor identities or shapes, this can silently produce wrong behavior. The fix is to separate the flatten step from the hook call in _run_warmup_forward so the hook always receives the original (unflattened) return value.

[timmoon10]

This is a good point.

[greptile-apps]

pre_warmup_hook/post_warmup_hook invocation count changed silently

In the original code these hooks were called once per callable, wrapping that callable's num_warmup_iters warmup iterations:

# original
for func_idx, func in zip(warmup_func_idx, warmup_func):
    if pre_warmup_hook is not None:
        pre_warmup_hook()
    for warmup_iter in range(num_warmup_iters):
        ...
    if post_warmup_hook is not None:
        post_warmup_hook()

The refactor moved the hooks outside both loops, so they now fire once globally, regardless of how many callables are registered. If a caller passes a pre_warmup_hook that relies on being invoked per-callable (e.g. to set up per-callable CUDA state before each callable's warmup), it will silently run only once and skip N-1 callables. This is an unintentional behavior regression for existing users of these hooks.

[timmoon10]

We should confirm that this doesn't cause problems, and probably update the docs for pre_warmup_hook/post_warmup_hook to be more explicit about its behavior. That said, I don't think this function will break the use-cases from https://github.com/NVIDIA/TransformerEngine/pull/2435/changes#r2850937378.

CC @lhb8125

[timmoon10]

Doesn't this sound weird? Let's make it "pre_forward" to be consistent with "pre_backward".

[timmoon10]

Why is capture_time_hooks[callable_idx]["forward_pre"] a dict? As far as I can tell, we just iterate through the hooks, so it would be more straightforward to make it a list (and similar for "forward", "pre_backward", "backward").

[timmoon10]

We should confirm that this doesn't cause problems, and probably update the docs for pre_warmup_hook/post_warmup_hook to be more explicit about its behavior. That said, I don't think this function will break the use-cases from https://github.com/NVIDIA/TransformerEngine/pull/2435/changes#r2850937378.

CC @lhb8125

[timmoon10]

I'd prefer if length mismatches triggered errors. This pattern is inviting silent bugs.

[timmoon10]

Nit: It's kind of uncomfortable that the APIs for the warmup and capture hooks are inconsistent. Also, the capture_time_hooks is slightly misleading because they are applied during warmup, not just during capture.

[timmoon10]

This is a good point.