@eellison, probably cuda graph tree related?

We figured out that torch.compiler.cudagraph_mark_step_begin() solved it.
The reason is simply that the graph is not released even though the value is detached. In eager, that wouldn't obviously be an issue.
Since in other cases you're explicitly asked to use torch.compiler.cudagraph_mark_step_begin when required, I assume it was always the case.
By the way, we've found another interesting issue which is that cudagraph with models that are not detached + mujoco rendering on-device causes a similar issue (very slow calls to the model), again without warning. This is again due to the fact that the graphs pile up in memory and conflict with mujoco until an OOM occurs. (We tried to come up with a simple minimal reprod but didn't manage to do so)

I'm mentioning this because in two unrelated cases we faced perf issues with "reduce-overhead" that had the same solution (torch.compiler.cudagraph_mark_step_begin()) but no warning and quite a different underlying cause.

I'm curious: what would an automated call to torch.compiler.cudagraph_mark_step_begin() cause when entering the compiled region, unless a certain flag is passed when calling compile? Am I correct that the usages that would (silently) break would be like

c1 = torch.compile(c1, mode="reduce-overhead")
c2 = torch.compile(c2, mode="reduce-overhead")

y = c1(x) # calls torch.compiler.cudagraph_mark_step_begin() automatically
z = c2(x, y) # calls torch.compiler.cudagraph_mark_step_begin() automatically
z.sum().backward() # does not propagate through c1 weights

The reason I'm asking is that if cudagraph conflicts with other libs (like mujoco) it might be hard to raise a warning consistently. Also, in a code like TDMPC2, we have several calls to cudagraph_mark_step_begin all across the codebase and it's a bit cumbersome to do that (esp givent that most users will not really understand what that is doing precisely).

One option could be

>>> c1 = torch.compile(c1, mode="reduce-overhead")
>>> y0 = c1(x)
>>> y1 = c1(x)
Warning: two consecutive calls of cudagraphed models have been made while the `free_tensor` argument has not been passed to the compiler. If you want to keep the graph in between iterations, please call compile with `free_tensor=False`. To silence this warning and keep the current behaviour pass `free_tensor=True`.

To state the current behavior:

There are two stages in cudagraphs, recording, and replaying. If you end a recording early it will manifest as an error, if you end a recording late it will manifest as being slow.

When we end a recording is controlled here:

I think what you are hitting is the final case - running_forwards_with_pending_backwards. Is the code you have above realistic ? When would you be running a forward for 1 million runs, accumulating a gradient, but not calling backward ? I'm surprised this is OOMing already. (or maybe I am misreading).

There is also this code that detects for when we should issue a warning:

I think the bug here is that for some reason the warning did not fire.

I'm curious: what would an automated call to torch.compiler.cudagraph_mark_step_begin()

We have this to some extent already, in that we detect when torch.compile is invoked and help use that to determine when to end a recording. But we are not willing to break the case of multiple, separate torch.compile wrappings of submodules in training. This is extremely common.

>>> c1 = torch.compile(c1, mode="reduce-overhead")
>>> y0 = c1(x)
>>> y1 = c1(x)
Warning: two consecutive calls of cudagraphed models have been made while the `free_tensor` argument has not been passed to the compiler. If you want to keep the graph in between iterations, please call compile with `free_tensor=False`. To silence this warning and keep the current behaviour pass `free_tensor=True`.

Today, the way this is handled in inference, is if you access y0 we'll give an informative error message. We did warning spam at one point but it tends to get ignored.

In training - we would expected y0 to participate in backward somehow.

We talked about potentially marking step in set_stance