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 pre_forward, 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. pre_forward 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, enabling per-callable hooks that execute outside the CUDA graph capture context (i.e., they are not recorded into the graph and will not be replayed). The motivation is Megatron-LM's FSDP integration, where parameter un-shard/re-shard allocations cannot occur during capture.

• Adds _canonicalize_capture_time_hooks and _run_capture_time_hooks helpers; hooks receive only the module as argument with signature hook(module) -> None.
• Refactors warmup into _run_warmup_forward / _run_warmup_backward closures to cleanly inject hook calls around each pass.
• Hook keys match _CAPTURE_TIME_HOOK_NAMES and the public docstring, though the PR description names them differently.

Confidence Score: 5/5

Safe to merge — hook dispatch, index aliasing, and is_training guards are consistent across both the _order and non-_order warmup/capture paths.

All four hook invocation sites correctly fire outside the CUDA graph capture context in both warmup and capture, and are properly guarded by is_training for the backward hooks. The callable_idx / per_callable_fwd_idx aliasing is consistent. The new test covers both _order=None and _order=[1,2,-2,-1] paths and directly asserts hooks do not fire during graph replay.

No files require special attention — the only finding is a naming inconsistency between the PR description and the actual key names used in the code.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Adds capture_time_hooks infrastructure. Hook dispatch, index aliasing (callable_idx vs. per_callable_fwd_idx), and both _order/non-_order paths look correct; backward hooks are correctly guarded by is_training in all paths.
tests/pytorch/test_cuda_graphs.py	New test exercises both _order=None and _order=[1,2,-2,-1] paths, verifies hooks fire outside CUDA capture context, and confirms hooks do not fire during graph replay.

Sequence Diagram

sequenceDiagram
    participant MGC as make_graphed_callables
    participant CTH as capture_time_hooks[i]
    participant Graph as CUDA Graph Context
    participant Module as Module

    Note over MGC: Warmup iterations
    MGC->>CTH: forward_pre_hooks[i](module)
    MGC->>Module: "func(*args) [outside graph]"
    MGC->>CTH: forward_hooks[i](module)
    MGC->>CTH: backward_pre_hooks[i](module)
    MGC->>Module: autograd.backward() [outside graph]
    MGC->>CTH: backward_hooks[i](module)

    Note over MGC: Capture phase
    MGC->>CTH: forward_pre_hooks[i](module)
    MGC->>Graph: enter capture
    MGC->>Module: "func(*args) [RECORDED]"
    MGC->>Graph: exit capture
    MGC->>CTH: forward_hooks[i](module)
    MGC->>CTH: backward_pre_hooks[i](module)
    MGC->>Graph: enter capture
    MGC->>Module: autograd.backward() [RECORDED]
    MGC->>Graph: exit capture
    MGC->>CTH: backward_hooks[i](module)

    Note over MGC: Replay — no hooks fire

Loading

_{Reviews (5): Last reviewed commit: "Initialize CUDA graph grad inputs" | Re-trigger Greptile}

[buptzyb]

Hi @ksivaman could you help review? Thanks!

[timmoon10]

We should add a test. make_graphed_callables has gotten too complicated for me to have any confidence of correctness just by looking at the code.

[timmoon10]

Our behavior with grad_inputs doesn't match the torch.nn.Module backward hooks. Here we pass in the grads from both inputs and weights, ignoring non-tensor inputs and tensor inputs without grads. torch.nn.Module passes in the grads from the inputs, including non-tensor inputs and inputs without tensor inputs with grads.

[buptzyb]

Addressed by adopting your later simplification suggestion. Capture-time hooks now have signature hook(module) and no longer expose grad_inputs or grad_outputs, so we no longer claim torch.nn.Module backward-hook semantics.

[timmoon10]

Having existence checks every time we access the hooks is quite burdensome. How about we preprocess capture_time_hooks so it has a fixed structure?

Function to preprocess `capture_time_hooks`

def _canonicalize_capture_time_hooks(
    num_callables: int,
    capture_time_hooks: Optional[List[Optional[Dict[str, Dict]]]],
) -> List[Dict[str, Dict]]:
    """Fill defaults in capture_time_hooks"""
    expected_keys = {
        "forward_pre_hooks",
        "forward_pre_hooks_with_kwargs",
        "forward_hooks",
        "forward_hooks_with_kwargs",
        "backward_pre_hooks",
        "backward_hooks",
    }

    # Make sure there is one entry per callable
    if capture_time_hooks is None:
        capture_time_hooks = [None] * num_callables
    if len(capture_time_hooks) != num_callables:
        raise ValueError(
            f"capture_time_hooks has {len(capture_time_hooks)} entries, "
            f"but there are {num_callables} callables."
        )

    # Make sure each callable has a dict with expected keys
    for callable_idx in range(num_callables):
        if capture_time_hooks[callable_idx] is None:
            capture_time_hooks[callable_idx] = {}
        hooks = capture_time_hooks[callable_idx]
        for key in expected_keys:
            if key not in callable_hooks:
                hooks[key] = {}
        if set(hooks.keys()) != expected_keys:
            raise ValueError(
                "Found unexpected keys in capture_time_hooks "
                f"({set(hooks.keys()) - expected_keys})."
            )

    return capture_time_hooks

Triggering pre-forward hooks

        # Call pre-forward capture-time hooks
        hooks = capture_time_hooks[callable_idx]["forward_pre_hooks"]
        hooks_with_kwargs = capture_time_hooks[callable_idx]["forward_pre_hooks_with_kwargs"]
        for hook_id, hook in hooks.items():
            hook_args = (
                (func, args, kwargs) if hook_id in hooks_with_kwargs else (func, args)
            )
            if hook(*hook_args) is not None:
                raise RuntimeError(
                    "capture_time_hooks forward_pre_hooks must not return a value "
                    "(args/kwargs must not be modified via hook return)"
                )

[buptzyb]

Implemented _canonicalize_capture_time_hooks based on this suggestion, with the reduced four-key hook contract. It validates one entry per callable, fills missing hook maps, and rejects unsupported keys. Invocation and return-value validation are centralized in _run_capture_time_hooks.

[timmoon10]

We are treating grad_outputs inconsistently between warmup and capture. During warmup, we pass in only the required grads. During capture, we pass in grads for all of the outputs (matching torch.nn.Module pre-backward hooks).

See @greptile-apps's review in #2831 (comment).

[greptile-apps]

Yes, exactly right — this is the bug I flagged. The fix in _run_warmup_backward is to construct grad_outputs with the same structure-preserving convention used in both capture paths, then strip Nones only for the actual torch.autograd.backward call:

grad_outputs = tuple(
    torch.empty_like(o) if o is not None and o.requires_grad else None
    for o in outputs
)

Then update the backward call to filter out Nones:

torch.autograd.backward(
    outputs_requiring_grad,
    grad_tensors=tuple(g for g in grad_outputs if g is not None),
)

This ensures backward_pre_hooks and backward_hooks receive a grad_output tuple that positionally aligns with the full outputs tuple (with None for non-grad outputs), exactly matching what static_grad_outputs provides in both capture paths and matching PyTorch's register_full_backward_pre_hook / register_full_backward_hook convention.

[buptzyb]

Now capture-time hooks no longer receive grad_outputs during either warmup or graph capture. So this problem disappears as well.

[timmoon10]

We put a lot of effort into matching the torch.nn.Module hook APIs, but they are causing correctness issues and we don't actually need them for our use-cases. Instead of making promises we can't fulfill, how about we simplify? We can just have "forward_pre_hooks", "forward_hooks", "backward_pre_hooks", "backward_hooks", each with signature hook(module).

[buptzyb]

Implemented as suggested. Thanks!