Hmm, we're specializing the types of the graph here without guaranteeing that they will be seen. This could lead to invalid peepholes in the forward of the differentiable graph. Could we try to rewrite this without removing the profiling nodes here ? The only property that we're guaranteeing stays consistent is the gradient of the inputs.

I think the flow we want is:

• CreateAutodiffSubgraphs

• Specialize the inputs of CreateAutodiffSubgraphs on their requires_grad property. This might require going into DifferentiableGraph and aggregaating the grad property based on uses.

• Add a Type Guard to the diff graph. In the fallback function, we probably want to inline the differentiable graph, since we are running it only if gradients have changed, and we may not need to differentiate at all.

• Next, when we differentiate the graph, it would be nice if we're not adding excess gradients, since we know which inputs require grad. (This doesn't have to be done in thisPR).

It is sort of annoying to have to keep the profile nodes in the graph, but I don't know if there's a better solution here.

I had the following in mind:

• CreateAutodiffSubgraphs
• Specialize the inputs of DifferentiableGraphs
• Add a guard based on the input requiresGrad.

My thinking was that the fallback function doesn't have a differentiable graph but only the nodes that have been moved inside the graph so it would be good to use as is.

Hmm, how does this work ? what do we want the graph to look like ?

if TypeCheck:
    DifferentiableGraph
else:
    FallBack(...)

within the Fallback, we probably don't want it to be a Differentiable Graph right, because the gradient properties are different? We might need to inline the DifferentiableGraph in the FallBack graph if i'm reading this right.

I was thinking that the fallback graph is inserted by insertTypeGuard using the subgraph attached to the guarded (ie the DifferentiableGraph) node. So it does not include the DifferentiableGraph itself. Then here the replaceFallbackGraphWithFallbackFunction replaces the fallback graph with the fallback function. (And that gets an executor so it can then do its own profiling + optimization.)

it's somewhat common for LSTMs (in a loop) to have tensor values whose requires_grad attributes flip from false (in the first iteration) to true (in the subsequent iterations). This behaviour would result in all_requires_grad_defined set to false. The other probable scenario would be when we switch from training to eval/test (i.e. true to false).

I think we could consider the following options:
a) enable re-profiling if we aren't sure, so we will profile a steady state (would cost us an extra call in a loop, to always go through a fallback)
b) gamble and assume that requires_gradients will be true. This strategy works out very nicely for both LSTMs and eval/train scenarios. We would take the fallback in the first iterations since requires_gradients would be false, but then we will be running the optimized path (no fallbacks). Similarly, we would spend more time in training than eval, so incurring the fallback penalty in the test/eval runs is preferable.
c) peel the first iteration, so profiler only sees requires_grad=true. We kinda tried it but very reluctantly.
d) do nothing and lose out on perf since we won't fuse in non-diff regions as long as even a single requires_grad=true anywhere in a graph.

Thoughts
@t-vi @eellison @jjsjann123 ?

Oha, indeed! I would go for b). If that doesn't work out, I'd probably go for actually querying the configuration in the ProfilingRecord for which is the most commonly seen.
This will apply to completely unseen things, too, though...

I also kinda like b). It seems to be a reasonable heuristic; we could always revise it later.

This will apply to completely unseen things, too, though...

that's true. we could check for unseen things as their types are just TensorType::get().

I'd probably go for actually querying the configuration in the ProfilingRecord for which is the most commonly seen.

I'm actually working on this in a separate PR. However, the initial version would aggregate across different invocations of a function rather than different iterations of a loop.

that's true. we could check for unseen things as their types are just TensorType::get().

I wonder whether we might just have a "seen" bit on the profiling node. I've been thinking about incremental profiling (where we leave profiling nodes for the bits we don't have records of and re-run the optimization once we hit them) and will probably play with it one of these days.

Yea, I think we should do b) but filter out TensorType::get() (unseen tensors). Incremental profiling could have value, i'm not sure, but is probably out of scope for this PRR

OK, I'll return to the taking out graphs with TensorType::get() for now and look into incremental profiling. Thank you for the insightful comments!

could you please add a comment here? I presume we need this this to minimize the number of nodes whose requires_grad isn't set because they aren't being profiled?

Is there a reason we can't just set requiresGrad attribute vs the complete type? TensorType::get().withRequiresGrad(n->type(attr::profiled_type)->requiresGrad())?

But then if withRequiresGrad is undefined (ie the case you mentioned), we would end with something indistinguishable from an unseen tensor. So I'd probably leave it as is for now.

OMG

Yeah. Eventually, I think the best fix would be to 1) re-profile the graph inside the differentiable graph (or do we do that already when using the differentiable graph operator?) or 2) apply the type annotations before creating autodiff graphs. So I understood @eellison that this isn't good to do at the moment, but I haven't understood well enough what is unsafe about it to see if that could be fixed.

1. is an easy enough - see switching from Code to GraphExecutor for gradient forward #47434.

2. yea is a little weird bc it gets away from our "fusers decide what to guard on" logic, and im also a little worried about proving what properties in the differentiable graph are guaranteed. For instance, if an Op takes in a NumberType, the output type could change if in one iteration the input is an Integer and the next a Float. I'm not saying that's particularly common, but it's just one example where strictly guarding on Tensor inputs does not guarantee consistent tensor types.

There's a test/jit/test_profiler, this test might be better there

Nice!

It would be nice if inputs without profiling information we just didn't guard on, instead of disabling differentiable graphs entirely, maybe as a follow up... I'm not sure this is acttually relevantt in practice though

I suppose with DifferentiableGraph's we're always fusing just a block, so all or no uses will have been executed

Hmm, when would this happen exactly? I'm a little confused here. I guess this is the LSTM first loop doesn't require grad ?

for LSTM in the first timestep the hidden state usually doesn't require gradients, the inputs would.

1. is an easy enough - see switching from Code to GraphExecutor for gradient forward #47434.

2. yea is a little weird bc it gets away from our "fusers decide what to guard on" logic, and im also a little worried about proving what properties in the differentiable graph are guaranteed. For instance, if an Op takes in a NumberType, the output type could change if in one iteration the input is an Integer and the next a Float. I'm not saying that's particularly common, but it's just one example where strictly guarding on Tensor inputs does not guarantee consistent tensor types.

Once we re-profile within differentiable graphs, we can remove the profile nodes before differentiating, and we not compute unnecesssary gradients. I'm not sure we're doing that now.

see:

(not for this PR, as a follow up)