autograd.Function supports vmap staticmethod

[zou3519]

Stack from ghstack:

• [pyfunctorch] Generate a more meaningful name for _SingleLevelAutogradFunction #90418
• Fix circular import in torch/autograd/function.py #90415
• functorch.jvp support for autograd.Function #90077
• Move private forward grad mode helpers to torch.autograd.forward_ad #90240
• -> autograd.Function supports vmap staticmethod #90037

This PR adds a vmap staticmethod to autograd.Function and a
corresponding vmap kernel for custom_function_call. These two items mean
that autograd.Function with a vmap staticmethod can be used with vmap.

class NumpyMul(torch.autograd.Function)
    staticmethod
    def forward(x, y):
        return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device)

    staticmethod
    def setup_context(ctx, outputs, x, y):
        ctx.save_for_backward(x, y)

    staticmethod
    def backward(ctx, grad_output):
        x, y = ctx.saved_tensors
        gx = None
        if isinstance(x, torch.Tensor) and x.requires_grad:
            gx = NumpyMul.apply(grad_output, y)
        gy = None
        if isinstance(y, torch.Tensor) and y.requires_grad:
            gy = NumpyMul.apply(grad_output, x)
        return gx, gy

    staticmethod
    def vmap(info, in_dims, x, y):
        x_bdim, y_bdim = in_dims
        x = x.movedim(x_bdim, -1) if x_bdim else x.unsqueeze(-1)
        y = y.movedim(y_bdim, -1) if y_bdim else y.unsqueeze(-1)
        result = NumpyMul.apply(x, y)
        result = result.movedim(-1, 0)
        return result, 0

API Spec

• the staticmethod takes two arguments (info, in_dims) as well as the
  unexpanded inputs (x, y).
• If we think about it as vmap(info, in_dims, *args), in_dims is a
  pytree with the same tree structure as args. It has None if the arg is
  not being vmapped over and an integer vmapped dimension index if it is.
• info is an object with metadata about the vmap. It currently has one
  field, info.batch_size. In the future we can extend this by adding
  things like the randomness information.
• If there is a single vmap going on, (x, y) are NOT BatchedTensors,
  they've already been unpacked.
• We expect the user to return a (outputs, out_dims) tuple. out_dims
  must "broadcast" to the same pytree structure as outputs.

Semantics

• vmap(NumpyMul.apply)(x) will apply the vmap staticmethod if there is
  one and will never actually run NumpyMul.forward.
• In order for the autograd.Function to support nested vmap (e.g.,
  vmap(vmap(NumpyMul.apply))(x), then the vmap staticmethod must call
  into operations that vmap understands (i.e. PyTorch operators or more
  autograd.Function).

At a high level, this PR:

• adds a vmap rule for custom_function_call

Testing

• Added some tests for in_dims and info
• Added vmap staticmethod to most of the autograd.Function in
  autograd_function_db and sent them through functorch's vmap-related
  OpInfo tests

Future

• Better error messages if the user gets the return contract wrong. I
  didn't include them in this PR because it might involve a refactor of
  some of the existing code in functorch/_src/vmap.py that will add
  ~200LOC to the PR, but LMK if you'd prefer it here.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/90037

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 Failures

As of commit 0826dca:

The following jobs have failed:

• linux-bionic-cuda11.6-py3.10-gcc7 / test (distributed, 2, 3, linux.8xlarge.nvidia.gpu)
• linux-bionic-cuda11.7-py3.10-gcc7 / test (distributed, 1, 3, linux.8xlarge.nvidia.gpu)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[zou3519]

Changes undone because it was better to unwrap dead wrappers before any interpreters run (see dispatch_functorch).

[zou3519]

sizes adjusted because we add extra dimensions of size 2 in vmap testing, so it is easier to debug when the sizes (before being pass through the vmap tests) don't have dimensions of size 2.

[soulitzer]

As a vmap noob who hasn't debugged many batching rules, I don't understand this lol, could you elaborate a little?

[zou3519]

Yup!

Let's say I incorrectly wrote the mul batching rule as:

def vmap(info, in_dims, x, y):
  return x * y

and then we tried:

x = torch.randn(2, 2)
y = torch.randn(2, 2, 2)
z = vmap(NumpyMul.apply, in_dims(0, 0))(x, y)

This would succeed and also be wrong -- we would end up with z = x * y, but the result should be z = x.unsqueeze(1) * y (we need to line up dimension 0 of x with dimension 0 of y). The test runner would fail with something like "output values are incorrect".

If the dimensions being vmapped over for x and y have a unique shape from the rest of the dimensions, then we get a loud error:

x = torch.randn(2, 4)
y = torch.randn(2, 3, 4)
z = vmap(NumpyMul.apply, in_dims(0, 0))(x, y)
# Error: can't broadcast Tensor of shape [2, 4] with Tensor of shape [2, 3, 4]

and with this error we can tell that we're doing something wrong (we forgot to change x into having shape [2, 1, 4] before multiplying it with y)

Now, how does this apply to the samples in the OpInfo?

• vmap testing takes the samples in the OpInfo and expands them to include an additional dimension of size 2.
• Then it just does vmap(op)(*expanded_inputs).
• Ideally the original samples in the OpInfo do not have dimensions of size 2 so we get a loud shape error instead of a "values are incorrect" error -- the shape error tells us we misaligned dimensions

[soulitzer]

Makes sense, thanks!

[samdow]

When is this None?

[zou3519]

_broadcast_to_and_flatten returns None if there is an error (if it's impossible to broadcast bdims to spec). I'm planning on writing better error messages in a followup

[samdow]

To check, this means that if a user's return looks like return (a, b), 0, this will imply that a and b are both batched at index 0

[zou3519]

Yes. It's the same rules as vmap's in_dims arg: vmap(f, in_dims=0)(a, b) implies both a and b are batched at index 0.

[samdow]

Looks good! Just a couple questions on the UX for returning batch dims

[soulitzer]

This is more of a pyfunctorch question, but is there a reason this plumbing cannot be done at the pyfunctorch level? If that is in place, all we need to do here is call autograd_function.vmap.

Analogously for autograd, interpreter.lower() would handling the boilerplate of unwrapping and unwrapping.

[zou3519]

This is more of a pyfunctorch question, but is there a reason this plumbing cannot be done at the pyfunctorch level? If that is in place, all we need to do here is call autograd_function.vmap.

That's a good question. I can see it going either way.

In C++, there are two types of batching rules that people write:
(1) unwrap(), lower(), do_something(), wrap()
(2) call a sequence of PyTorch operations

Most of our C++ batching rules are written as the former. CompositeImplicitAutograd and some other things are written as the latter (the batching rule for cross_entropy_loss calls log_softmax and nll_loss without unwrapping and lowering) There are a non-zero amount of rules that are a hybrid between the two (yikes).

pyfunctorch is supposed to be consistent with our C++ dispatcher, so it should in theory let people (developers) do either. If we did this plumbing at the pydispatcher level, this would force (1) onto every PyOperator, right? For autograd.Function with vmap staticmethod, we've made the decision that we are forcing (1) onto people for simplicity (so they don't need to manually call interpreter.lower() and understand what an interpreter is).

[soulitzer]

Interesting, I didn't think about the case where we want to register a composite. But one workaround could be: if we force (1) on people, maybe we could replace that functionality with an API specifically for registering decompositions for layers.

Sounds like this might prevent people from doing the hybrid approach, but maybe we want to discourage those anyway, and this would actually be the way of enforcing that?

[zou3519]

Hybrid approach isn't bad if there's a good reason for it (which there is in the C++ code -- it's difficult to write a batching rule when your operation does broadcasting).

For this PR I'll land it as-is, but we can revisit the situation in the future.

Analogously for autograd, interpreter.lower() would handling the boilerplate of unwrapping and unwrapping.

To check my understanding of this comment -- would one no longer have an interpreter.lower() inside the custom_function_call_grad function if we gave the same treatment to the grad transform?

[soulitzer]

To check my understanding of this comment -- would one no longer have an interpreter.lower() inside the custom_function_call_grad function if we gave the same treatment to the grad transform?

Oh what I meant was: In c++ functorch, the "send to next" of autograd interpreter handles unwrapping (which is pop layer off stack + unwrap), but in pyfunctorch the "send to next" for autograd interpreter is just lower (just pop layer of stack), so that is why user has to manually unwrap. I don't know if we want to do anything about that though.

[soulitzer]

I feel like if we implemented vmap for this one, it would be composable in at least one way: grad(vmap. vmap(grad probably won't work still.

[soulitzer]

(just a comment, not really actionable unless you also wanted to test this case)

[zou3519]

I don't actually know how to implement vmap for this one (and this is intentional -- I don't think users should be trying to implement vmap support for it). NumpyCubeNotComposableAutogradFunction returns a (tensor, ndarray). One is only allowed to specify bdims (dimension being vmapped over) for Tensors (and the only return from vmap can be Tensors), but we have an ndarray here

[soulitzer]

How is this different from the get_unwrapped function that we used previously? Do we want to update those uses to use unwrapBatched?

[zou3519]

Are you talking about functorch._C.get_unwrapped? Given a {Tensor, BatchedTensor, TensorWrapper, FunctionalTensorWrapper}, get_unwrapped returns the underlying Tensor (if applicable).

Given a {Tensor, BatchedTensor} and an integer level, unwrapBatched returns the underlying Tensor (if applicable) and the dimension it is being vmapped over.

[soulitzer]

As a vmap noob who hasn't debugged many batching rules, I don't understand this lol, could you elaborate a little?

[soulitzer]

LGTM, just had a couple small questions

[facebook-github-bot]

This pull request has been merged in 3049d99.