Training with custom, quantized datatype

username: ashari4
cooked: <p>I want to train a model that has weights in a custom, quantized datatype (&lt; 16 bits) to bring down memory usage. To avoid precision loss, I still want the gradient to be computed in torch.bfloat16. A torch.float32 copy of the weight would be used for accumulation.</p>
<pre><code class="lang-auto">import torch
class QuantizedModuleFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, w_quantized):
        # will eventually be some actual computation
        return torch.zeros(128,128, dtype=input.dtype, device=input.device)

    @staticmethod
    def backward(ctx, grad_output):
        # will eventually be some actual computation
        grad_input = torch.zeros(128,128, dtype=grad_output.dtype, device=grad_output.device)
        grad_w =  torch.zeros(128,128, dtype=torch.bfloat16, device=grad_output.device)
        return grad_input, grad_w

class QuantizedModule(torch.nn.Module):
    def __init__(self):
        super(QuantizedModule, self).__init__()
        # will eventually be torch.&lt;some custom dtype&gt;
        # See question 1
        self.w_quantized = torch.nn.Parameter(torch.zeros(128, 128, dtype=torch.complex64))
        self.w_fp32 = torch.zeros(128, 128, requires_grad=False)

    def forward(self, input):
        return QuantizedModuleFunction.apply(input, self.w_quantized)

def custom_optimizer(w_quantized, w_fp32, grad):
    """ will eventually update w_fp32 and w_quantized and zero out grad"""
    pass

def test_training_loop():
    input = torch.zeros((128,128), requires_grad=True, dtype=torch.bfloat16)
    model = QuantizedModule()
    for i in range(10):
      ans = model(input)
      loss = ans.sum()
      loss.backward()
      if model.w_quantized.grad.dtype != torch.bfloat16:
          # see question 2
          print(f"gradient dtype expected to be torch.bfloat16 but is instead {model.w_quantized.grad.dtype}")
      if (i+1) % 2 == 0:
        # update weights every 2 microbatches
        custom_optimizer(model.w_quantized, model.w_fp32, model.w_quantized.grad)
</code></pre>
<ol>
<li>Currently, I am using <code>torch.complex64</code> for the weight; but I would like to use a custom, quantized datatype. Are there any suggestions on how to add python bindings for a new datatype?</li>
<li>Despite the backward function returning <code>torch.bfloat16</code>, <code>model.w_quantized.grad.dtype</code> is <code>torch.complex64</code> because an implicit cast is performed here: <a href="https://github.com/pytorch/pytorch/blob/d03d43df527e48771875537ad20212d5cb333215/torch/csrc/autograd/engine.cpp#L704-L706" class="inline-onebox" rel="noopener nofollow ugc">pytorch/engine.cpp at d03d43df527e48771875537ad20212d5cb333215 · pytorch/pytorch · GitHub</a>. Is there a preferred method for computing the gradient in a different type than the weight? One option is to add a flag (e.g <code>requires_grad_diff_type</code>) to the <code>TensorOptions</code> object.</li>
</ol>
<p>Alternatives I’ve considered:</p>
<ul>
<li>An alternative is to define the model weight in <code>torch.bfloat16</code> and have the forward pass internally create a quantized copy. This is less desirable since it increases memory usage; there would be both a <code>bfloat16</code> and a quantized copy in memory. Additionally, the quantized copy would be recreated every microbatch even though the weights only change every 2 microbatches.</li>
<li>I took a look at the torch quantization feature <a href="https://pytorch.org/docs/stable/quantization.html" rel="noopener nofollow ugc">Quantization — PyTorch 1.11.0 documentation</a>. But in my usage scenario, I want the weights to be trained in the quantized datatype, which doesn’t seem possible.</li>
</ul>

username: albanD
cooked: <p>Hey!</p>
<ol>
<li>Adding a new native dtype is pretty tricky and not really possible today I’m afraid. You can use a custom subclass though to do something like that: <a href="https://dev-discuss.pytorch.org/t/what-and-why-is-torch-dispatch/557" class="inline-onebox">What (and Why) is __torch_dispatch__? - frontend API - PyTorch Dev Discussions</a> and a collection of examples <a href="https://github.com/albanD/subclass_zoo" class="inline-onebox">GitHub - albanD/subclass_zoo</a>
</li>
<li>The constraint in the autograd engine is mostly because cross dtype tensor/grad is not something that is used today and that is a common error when implementing backwards. So it is safer to have this check. We are open to lifting this check for particular Tensor types (in particular subclasses that I mentioned above).</li>
</ol>

username: ashari4
cooked: <aside class="quote no-group" data-username="albanD" data-post="2" data-topic="152132">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/alband/48/215_2.png" class="avatar"> albanD:</div>
<blockquote>
<p>We are open to lifting this check for particular Tensor types (in particular subclasses that I mentioned above).</p>
</blockquote>
</aside>
<p>Thanks <a class="mention" href="/u/alband">@albanD</a> . Assuming I create a python subclass, how would the python type be retrieved from  here: <a href="https://github.com/pytorch/pytorch/blob/d03d43df527e48771875537ad20212d5cb333215/torch/csrc/autograd/engine.cpp#L704-L706" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/d03d43df527e48771875537ad20212d5cb333215/torch/csrc/autograd/engine.cpp#L704-L706</a></p>

username: albanD
cooked: <p>Today, this dtype will be the “true” dtype of the subclass you created.<br>
It will be either the dtype of the Tensor you gave to <code>make_subclass</code> or the dtype value passed to <code>make_wrapper_subclass</code>.</p>
<p>We are looking into ways to extend this, but since c++ is strictly typed, it is not very easy to do so.<br>
But such a check can be bypassed like the one below when <code>tensor.is_tensor_subclass()</code>.</p>

username: ashari4
cooked: <p><a class="mention" href="/u/alband">@albanD</a> Thanks for your suggestions. I’ve created <a href="https://github.com/albanD/subclass_zoo/pull/39" rel="noopener nofollow ugc">[WIP] quantized tensor by ashari4 · Pull Request #39 · albanD/subclass_zoo (github.com)</a>, with some questions. I would appreciate a review.</p>