make cudagraphs a transform

[t-vi]

This introduces thunder.transform.cudagraph.CUDAGraphTransform to replace use_cudagraphs=True.

For more detailed control, users can subclass the transform and e.g. override can_fuse.

[nikitaved]

we are missing num_static_inputs here for a reason? If we get several non-isomorphic graphs, we should probably think about how to handle this parameter better, maybe through a callback. But that's not relevant now...

[t-vi]

I do think that this is one of the things where the default transform is not ideal, but it is not something that was used before: the parameter was there with a default, but there was no way of providing it.
(And I am not sure that it is correct to have a nontrivial trace-global parameter for it either.)

[nikitaved]

It is used in the backward pass. We are not loosing it there, are we?

[t-vi]

Ouch, right.

[nikitaved]

But yes, the design of this parameter is so-so as it was not expected to have had graph breaks...

[t-vi]

but which were the args covered by it?

[nikitaved]

So, this is the removed code:

            if cd.use_cudagraphs:
                from thunder.executors.cudagraphex import cudagraphex

                computation_trc = cudagraphex.fusion_pass(computation_trc)
                computation_traces.append(computation_trc)

                if backward_trc is not None:
                    backward_trc = cudagraphex.fusion_pass(backward_trc, num_static_inputs=len(backward_trc.args[0][0]))
                    backward_traces.append(backward_trc)

[t-vi]

So it seems that these were the saved for backwards of (???). Is the assumption here that they were either static in the forward (parameters) or copied to the input area of the forward cuda graph?

But with the old code:

import torch, thunder
with torch.device("cuda"):
    m = torch.nn.Linear(2, 3)
    inp = torch.randn(1, 2, requires_grad=True)
jm = thunder.jit(m, use_cudagraphs=True)

res = jm(inp)
grads = torch.autograd.grad(res.sum(), (inp, *m.parameters()))

the forward has no cudagraph, so in that case, having the input as a fixed is not really correct. (admittedly, a corner case).

[t-vi]

I guess the happy part is that we will get the parameters as static anyways and we will have to look into the buffers for our own good...

[crcrpar]

would there be a plan to add a test which demonstrates the composability with e.g. fsdp?

[t-vi]

would there be a plan to add a test which demonstrates the composability with e.g. fsdp?

That would be awesome, I really need to get the env for running distributed tests back. :(

[t-vi]

So as we talk about plans:

• refine the static inputs a bit,
• refine the caching, in particular don't use a global cache, but a per-transform one, to let things go out of scope,
• refine the operator selection for what is cuda-graphable,
• check composibility - I'm guessing the communication prims don't mix with cudagraphs?

To my mind, this greatly improves our starting position (similar to but smaller than Nik's work to make this a (then hardcoded) transform instead of a callable wrapper).

[mruberry]

This is cool; how do you see cuda graphs as a transform? If we conceptually separate the transforms like this:

pre-execution transforms
execution transform
post-execution transforms
"destructive" transforms

Where would the cuda graphs transform go? If it's a post-execution (or destructive) transform, should we think about how executors might label their operations as being cuda graphs compatible (or not)? If it's a pre-execution transform, then does it preclude using certain executors in the execution transform?

[t-vi]

So it is after transform for execution ("post optimization") but our experimentation (which informs some of the CUDAGraph work) caused us to go before the "del" pass.
I think it would be cool if symbols in general had a tag whether they are suitable for CUDAGraph inclusion. The current thing has the advantage that if people have particular ideas about this, they can just subclass and override can_fuse.
Also, CUDAGraphs could enormously benefit by more information about the memory effects of operators, but that is in the future.

[lantiga]

👍

[t-vi]

I have merged this after coordination with @IvanYashchuk , we will address further review comments in a follow-up.