torch.compiler.cudagraph_mark_step_begin

torch.compiler.cudagraph_mark_step_begin()

Indicates that a new iteration of inference or training is about to begin.

CUDA Graphs will free tensors of a prior iteration. A new iteration is started on each invocation of torch.compile, so long as there is not a pending backward that has not been called.

If that heuristic is wrong, such as in the following example, manually mark it with this api.

@torch.compile(mode="reduce-overhead")
def rand_foo():
    return torch.rand([4], device="cuda")

for _ in range(5):
    torch.compiler.cudagraph_mark_step_begin()
    rand_foo() + rand_foo()

For more details, see torch.compiler_cudagraph_trees