Add tutorial for trainable tensor subclass

[andrewor14]

Summary: The new tutorial provides an example of how to implement a trainable tensor subclass that wraps quantized data. This extends the existing MyDTypeTensor with a few necessary steps to ensure proper gradient updates, namely:

1. Define a differentiable constructor
2. Define backward pass for ops of interest (e.g. torch.nn.functional.linear)
3. Handle special ops used by the optimizer (e.g. aten.add, aten.add_)

Test Plan:
python tutorials/developer_api_guide/my_trainable_tensor_subclass.py

 * step 0: weight grad = tensor([-0.3854,  0.3851,  0.2130], device='cuda:0'), weight value = tensor([  41,  -73, -110], device='cuda:0', dtype=torch.int8)
 * step 1: weight grad = tensor([1.3488, 0.1335, 0.8074], device='cuda:0'), weight value = tensor([-107,  107,  116], device='cuda:0', dtype=torch.int8)
 * step 2: weight grad = tensor([-0.1505,  0.0958, -0.2342], device='cuda:0'), weight value = tensor([114, -98,  94], device='cuda:0', dtype=torch.int8)
 * step 3: weight grad = tensor([-0.4069,  0.7332,  0.1454], device='cuda:0'), weight value = tensor([-101,   54,  111], device='cuda:0', dtype=torch.int8)
 * step 4: weight grad = tensor([-0.0226,  0.0785, -0.0824], device='cuda:0'), weight value = tensor([  48,   31, -109], device='cuda:0', dtype=torch.int8)
 * step 5: weight grad = tensor([-2.0708,  0.6241, -2.0108], device='cuda:0'), weight value = tensor([112, -16,   5], device='cuda:0', dtype=torch.int8)
 * step 6: weight grad = tensor([-0.6014,  2.3908,  0.3403], device='cuda:0'), weight value = tensor([-17,  28,  69], device='cuda:0', dtype=torch.int8)
 * step 7: weight grad = tensor([ 1.1537, -2.1295,  0.5759], device='cuda:0'), weight value = tensor([-36, 116, 101], device='cuda:0', dtype=torch.int8)
 * step 8: weight grad = tensor([ 0.4791,  0.9322, -0.9088], device='cuda:0'), weight value = tensor([ 25, -24,  75], device='cuda:0', dtype=torch.int8)
 * step 9: weight grad = tensor([-0.4855,  0.1751, -0.4969], device='cuda:0'), weight value = tensor([ 26, 111,   1], device='cuda:0', dtype=torch.int8)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 20f3dc0 with merge base 53b6b78 ():
💚 Looks good so far! There are no failures yet. 💚

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

[jerryzh168]

are add_/add both going to be differentiable as well?

[andrewor14]

I think they don't need to be. These are just used by the optimizer

[jerryzh168]

thanks, this looks great!

[andrewor14]

Ok, merging this. Thanks!

[jerryzh168]

sorry I guess I didn't mean to dequantize this, but was asking if this call into something like F.linear(grad_output, weight_tensor) and be dispatched to the quantized linear impl of weight_tensor.

[andrewor14]

oh actually the old code was also doing dequantize math. I just rewrote it so it's cleaner. Not sure if I understand how we can dispatch to the quantized linear impl, since this is called from F.linear(..., weight_subclass_tensor) already?