Multi-output derivative formulas can save unnecessary tensors

[soulitzer]

Usually when you have a derivative formula, e.g. mm:

- name: mm(Tensor self, Tensor mat2) -> Tensor
  self: mm_mat1_backward(grad, mat2, self.sym_sizes(), self.sym_strides(), self.layout(), 1)
  mat2: mm_mat2_backward(grad, self, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), 1)
  result: at::mm(self_t, mat2_p) + at::mm(self_p, mat2_t)

The structure informs autograd codegen what tensors need to be saved to compute which gradients, i.e. if only self requires grad, I only need to save mat2, and if only mat2 requires grad, I only need to save self.

In VariableType, the following logic is generated for mm:

if (_any_requires_grad) {
    grad_fn = std::shared_ptr<MmBackward0>(new MmBackward0(), deleteNode);
    grad_fn->set_next_edges(collect_next_edges( self, mat2 ));
    if (grad_fn->should_compute_output(1)) {
      grad_fn->self_ = SavedVariable(self, false);
    }
  ...
    if (grad_fn->should_compute_output(0)) {
      grad_fn->mat2_ = SavedVariable(mat2, false);
    }
}

However, when you have a single derivative formula that produces multiple outputs, autograd codegen no longer has visibility into what tensors need to be saved in order to compute which gradients.

- name: matmul(Tensor self, Tensor other) -> Tensor
  self, other: matmul_backward(grad, self, other, grad_input_mask)

This yields the following unoptimized logic in autograd kernel:

 if (_any_requires_grad) {
    grad_fn = std::shared_ptr<MatmulBackward0>(new MatmulBackward0(), deleteNode);
    grad_fn->set_next_edges(collect_next_edges( self, other ));
    grad_fn->self_ = SavedVariable(self, false);
    grad_fn->other_ = SavedVariable(other, false);
  }

One common case where this can matter is if I'm doing some fine-tuning, e.g. I want to requires_grad=False for the parameters of a bunch of layers in the middle of a network.

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @lezcano @Varal7 @cpuhrsch @jbschlosser @bhosmer @drisspg @gchanan

[zou3519]

How would you fix this? By splitting single derivative formulas into multiple?

[albanD]

I don't think there is a way for us to do anything here: if you can split up the formula, you can do so.
But if computing each gradient separately is more expensive, then we could imagine a custom rule where you can provide both separate and merged formulas. But that would be quite a bit of work and I don't think we see a significant benefit for it?

Also in the great new world of torch.compile, the un-used computations will be DCEd away :D

[soulitzer]

I haven't checked too many other ones, but matmul autograd explicit kernel (specific to nested) seems splittable

[lezcano]

Sometimes, the issue with splitting them is that, if you have to implement the second derivatives, you need to implement the cross derivatives twice (and these will be recomputed if the user asks for these) which is a bit annoying.

[lezcano]

of course, this is not the case of matmul

[ezyang]

Another crappy fix is to do something like:

- name: matmul(Tensor self, Tensor other) -> Tensor
  self, other: matmul_backward(grad, self, other, grad_input_mask)
  self: blah blah
  other: blah blah

[albanD]

Yeah that's what I hinted at above, but we should only do that if we have cases with clear benefit. And given the new compile world, I don't think it will happen.

[soulitzer]

@nikitaved has a fix for this now in #103750 which introduces a macro to hint autograd codegen at what to save.
What's left to do is update existing formulas to use this.