Repository: pytorch
Issue: custom gradient for int8
Labels: oncall: quantization

### 🚀 The feature, motivation and pitch

**Feature motivation:**
[Default pyTorch quantization aware training](https://pytorch.org/docs/stable/quantization.html) uses "fake-quantization" approach. For example for int8 quantization we will use float type during training and int8 during inference. There several disadvantages in this approach:
* It does not follow approach of "what you train is what your serve", as a result there can be numerical difference between training and inference (it emulates int ops using float ops and as soon as float summation goes beyond of 23 bits of float mantissa there will be a numerical difference after converting such model to int8).
* During training, it uses float operations, but it can use int8 operations to speed up training too.
* Current fake quantization aware training api uses both [python](https://github.com/pytorch/pytorch/tree/bb0f3df56279dff6782767804caebe5f6ffb2439/torch/ao/quantization) and [c++](https://github.com/pytorch/pytorch/tree/bb0f3df56279dff6782767804caebe5f6ffb2439/torch/csrc/jit/passes/quantization) levels, which makes it harder to support and extend. Instead if we use native pytorch types (e.g. int8, fp8) we could simplify quantization lib design and specify model quantization all in python. For example fp8 type is already used in [float8_experimental](https://github.com/pytorch-labs/float8_experimental).

To train a model with native int8 type I need two things:
* Control over output type in a matmul, e.g. for input int8 type return int32, as described in the [issues/128283](https://github.com/pytorch/pytorch/issues/128283)
* Custom gradient over int8 casting.

**Current issue:**
int8 tensors does nor have a gradient as discussed in ([link1](https://discuss.pytorch.org/t/is-the-type-conversion-differentiable/45436); [link2](https://discuss.pytorch.org/t/keep-the-gradient-function-when-changing-the-data-type-of-tensor/161639)). But it would be great if customer can specify their own gradient for int types in a straight  through estimator(STE).

There are two options to to do straight  through estimation (as shown below):
* Use custom backward (ste_custom_bwd).
* Use detach (ste_detach).

```
import torch

def ste_custom_bwd(func):
    """Straight through estimator based on custom backward."""
    class StraightThroughEstimator(torch.autograd.Function):
        @staticmethod
        def forward(ctx, input):
            return func(input)
        
        @staticmethod
        def backward(ctx, grad_output):
            return grad_output

    return StraightThroughEstimator.apply

def ste_detach(input, func):
  """Straight through estimator based on detach."""
  return input - input.detach() + func(input).detach()

def cast_int8(x):
    return x.to(torch.int8)
```

STE based on detach works ok:
```
x = torch.tensor([2.0], requires_grad=True)
y = ste_detach(x, cast_int8)
y.backward()
print("grad" + str(x.grad))
```

But STE based on custom backward pass, fails with expected message "element 0 of tensors does not require grad and does not have a grad_fn". 
```
x = torch.tensor([2.0], requires_grad=True)
y = ste_custom_bwd(cast_int8)(x)
y.backward()
print("grad" + str(x.grad))
```

It is not a blocker, because STE based on detach works, but it would be great if customer can specify custom backward pass for any type, for consistency with the API.


### Alternatives

_No response_

### Additional context

_No response_

cc @jerryzh168 @jianyuh @raghuramank100 @jamesr66a @vkuzo @jgong5 @Xia-Weiwen @leslie-fang-intel @msaroufim