Add test util for basic tensor subclass functionalities

[jerryzh168]

Summary:
This is a small starting point for testing low precision tensor subclass functionalities we can add more test cases for training, tensor parallel, FSDP in the future

right now it tests:

• tensor flatten/unflatten
• constructing low precision tensor with different device/dtype
• move tensor subclass from device1 to device2
• transpose works
• linear works (weight only quantization with the low precision tensor)

It can be extended with new tensor subclasses or test cases by overriding the class variables: e.g.

class MyTensorSubclassTest(TorchAOBasicTestCase):
    TENSOR_SUBCLASS = LUTQuantizedTensor
    FACTORY_FN = to_lut_quantized_intx
    kwargs = {
        "target_dtype": torch.uint8,
    }
    # minimum sqnr for linear operation when the weight is quantized to low precision
    # with the above setting
    LINEAR_MIN_SQNR = 40

Test Plan:
python test/utils.py

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 75d9c90 with merge base e05635e ():
💚 Looks good so far! There are no failures yet. 💚

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

[bdhirsh]

one kind of annoying thing that makes __tensor_flatten__ and friends difficult to test in isolation is that we want their implementations to be side-effect-free / idempotent. This is kind of difficult to test by calling the fn directly, although usually you will find out if your implementation is wrong pretty quickly because usages of compile with your subclass will break.

This test suite seems like a good start (if the goal is "a test suite for all AO-specific quantization subclasses that conform to a specific API"). Do you want / plan to add more compile-related tests?

[jerryzh168]

I see, yeah I can add more compile related tests, although I'm not exactly sure what are all the scenarios I should be testing, I can also start with simple cases like just compile the quantized weights and run on linear op as well, but I'd like to hear what's your thoughts on rough categories of compile related tests we can add as well

[bdhirsh]

A few that come to mind are:

(1) simple test where inputs + outputs of the function f are your subclass and you do some basic compute (return inp + inp)
(2) subclass constructor in the graph (e.g. inputs to f are plain tensors, output is a subclass)

[jerryzh168]

sg, I can add these in a separate PR

[bdhirsh]

seems fine!