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

[buptzyb]

Description

Summary

This PR adds capture_time_hooks to make_graphed_callables for per-callable work that must run during warmup and CUDA Graph construction but remain outside the CUDA Graph capture context. This is needed for non-capturable operations such as CPU-side state updates and FSDP parameter un-shard/re-shard work that may allocate memory.

These callbacks are graph-construction hooks, not PyTorch Module hooks: they receive only the original callable/module, must return None, are not recorded into the graph, and are not replayed.

Changes

• Add one capture_time_hooks entry per callable, with support for forward_pre_hooks, forward_hooks, backward_pre_hooks, and backward_hooks.
• Run the hooks around each callable's forward and backward during warmup and graph construction, outside the corresponding CUDA Graph capture context.
• Support both capture paths: no-order warmup runs all forwards followed by backwards in reverse order, while ordered warmup follows _order across chunks and microbatches. Existing backward_dw() warmup behavior is unchanged.
• Call the existing pre_warmup_hook and post_warmup_hook once around the complete warmup phase rather than once per callable.
• Reject modules with registered backward pre-hooks consistently with the existing checks for other pre-registered Module hooks.
• Add parameterized CUDA coverage for both _order=None and _order=[1, 2, -2, -1].

Hook contract

capture_time_hooks is a list with one entry per callable. Each entry may be None or a dictionary containing any subset of the four supported hook groups. Each group is a {hook_id: hook_fn} dictionary, and every hook has signature hook(module) -> None.

The list length must match the number of callables. Missing hook groups are treated as empty, unsupported group names raise ValueError, and a non-None hook return raises RuntimeError. The hooks cannot replace inputs, outputs, grad_inputs, or grad_outputs.

Normal PyTorch Module hooks may still be registered after graph construction.

Execution order

For each callable, the order is:

1. forward_pre_hooks
2. forward execution
3. forward_hooks
4. backward_pre_hooks (training only)
5. backward execution (training only)
6. backward_hooks (training only)

During graph construction, pre-hooks run before entering the corresponding capture context and post-hooks run after leaving it.

Tests and validation

The CUDA test verifies:

• exact forward/backward hook order;
• original module identity;
• execution outside CUDA Graph capture;
• invocation during every warmup iteration and graph construction;
• no invocation during graph replay.

Validation performed:

• focused capture-time hook test on H100: 2 passed;
• full tests/pytorch/test_cuda_graphs.py on H100: 412 passed, 423 skipped;
• pylint==3.3.1, pre-commit, Python compilation, and diff checks passed.

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!