Use functionalization to achieve good coverage on vmap(vjp(f)) compositions

[zou3519]

The backward formula for operators like diag call in-place operations. We cannot vmap over these so we want to develop some way to functionalize the backward formula. However, "conditional functionalization" (#235) right now only works on operators: we can say, "given this native_functions operator, we will apply the functionalization pass to it".

The goal of this issue is to figure out how to use functionalization to functionalize the backward pass of certain operations when vmap is active.

Approach 1

Just register the backward formula as a composite operator. This will be useful for things in the future like dynamic shape tracing where we need the backward formula to be available in Python.

Tradeoffs: There is resistance to making backward formulas into operators because (1) it increases the size of our operator library and (2) the native_functions operator library is potentially a FC/BC surface. Perhaps we would have to couple this approach with some BE improvements like being able to mark a native_functions operator as "never serializable"?

Approach 2

Use Autograd hooks. For Autograd hooks to work, we would need to:

1. be able to enable functionalization before the autograd node gets called
2. be able to unconditionally disable functionalization after the autograd node is done executing. This needs to happen even if an error happens during the backward pass.

It's unclear how much work this approach is.

Test Cases

Here are some test cases for where we would like to conditionally functionalize the backward pass of these operators:

• torch.masked_select
• torch.clamp
• torch.diag
• torch.index_select
• torch.topk
• torch.prod
• torch.mode
• torch.sort

cc @albanD @bdhirsh for discussion

[gchanan]

this means, you want to test the functionalization pass on AOT grad tracing of these operators?

[zou3519]

this means, you want to test the functionalization pass on AOT grad tracing of these operators?

No, we want to use the functionalization pass to get vmap coverage for the backward pass of these operators.

EDIT: for some more clarification:

• the backward formula for operators like diag call in-place operations.
• We cannot vmap over these so we want to develop some way to functionalize the backward formula.
• However, "conditional functionalization" right now only works on operators: we can say, "given this native_functions operator, we will apply the functionalization pass to it"
• Many backward formulas are not operators and there is resistance to making backward formulas into operators because (1) it increases the size of our operator library and (2) the native_functions operator library is potentially a FC/BC surface. So we want to investigate finding a different way to functionalize the backward formula.

EDIT: I updated the body of this issue

[bdhirsh]

I'm not super familiar with autograd hooks. But would a cheap version of option 2 look equivalent to adding that conditional logic into each of the existing backward formula? E.g. in FunctionsManual.cpp:

// The current version
Tensor clamp_backward_old(const Tensor & grad, const Tensor &self, const optional<Scalar> & min, const optional<Scalar> & max) {
  // clamp: gradients not defined on min and max, so we return the subgradient 1 for these cases.
  if (max && min) {
    auto zero = at::scalar_tensor(0., grad.options());
    return where((self >= *min).logical_and_(self <= *max), grad, zero);
  } else if (min) {
    auto zero = at::scalar_tensor(0., grad.options());
    return where(self >= *min, grad, zero);
  } else if (max) {
    auto zero = at::scalar_tensor(0., grad.options());
    return where(self <= *max, grad, zero);
  } else {
    return grad;
  }
}

// The updated version
Tensor clamp_backward(const Tensor & grad, const Tensor &self, const optional<Scalar> & min, const optional<Scalar> & max) {
  if (c10::impl::tls_local_dispatch_key_set().included_.has(c10::DispatchKey::FuncTorchBatched)) {
    c10::impl::IncludeDispatchKeyGuard include_guard(c10::DispatchKeySet(c10::DispatchKey::Functionalize));
    return clamp_backward_old(grad, self, min, max);
  }
  return clamp_backward_old(grad, self, min, max);
}

I'm only asking because if I understand option 2 correctly, it has similar downsides to what I put above: we need to special case vmap in some way in core, for a hardcoded list of ops (the above code is probably a little more boilerplate heavy, but we could maybe template some of it away).

I kind of like approach 1 more, although I think I have a pretty loose understanding of to what extent our compatibility concerns come from serialization vs. other stuff. If we teach the codegen how to make an op "not serializable" (I'm not sure how to do this), don't we still have to worry about breaking people who call the op directly? Or I guess the idea is that nobody should be calling backward ops directly anyway, so we only really have to deal with serialized models that call the ops during autograd.

[albanD]

I'm afraid that at this point, with torch dispatch, all the functions in native functions are public API (even if not bound to python). So moving these formulas effectively make them public API and prevents us from changing them.

[bdhirsh]

Maybe we need a native_functions.yaml config for "dispatcher aware but not exposed to python and not usable with the Python key" 😛 (I'm only half joking)

[zou3519]

Hmm, one solution that came up in some discussion with Alban is when vmap sees that we are executing a backward pass, it could decide to functionalize the entire backward pass. For functionalization we're worried about two things w.r.t. performance:

1. overhead
2. additional copies

Now the question is, how bad is the performance overhead for functionalization, really? I'm not too worried about the wrapping/unwrap in FunctionalTensorWrapper. The copies seem scary,

but maybe the additional copies aren't too bad? Additional copies generally come from calling in-place operations after clone. For example I've seen code like the following

return x.clone().fill_(1)

which would functionalize into:

# I don't think an out_of_place_fill exists yet
x.clone().out_of_place_fill(1)

and functionalization isn't smart enough to remove the intermediate clone() yet. But maybe this happens infrequently enough that we shouldn't worry about it

[zou3519]

But would a cheap version of option 2 look equivalent to adding that conditional logic into each of the existing backward formula? E.g. in FunctionsManual.cpp:

@bdhirsh yes, a cheap version of option 2 would look like adding the conditional logic into each of the existing backward formulas.