overwrite to() for QTensor and QBitsTensor

[SunMarc]

What does this PR do ?

This PR fixes the to() method so that it is applied to the tensors _data, _scale and _zeropoint when using QTensor or QBitsTensor. Before this PR, doing the following QLinear().to(device) would not change the device of these tensors.

I ran the following tests (that were not passing before) : python -m pytest -sv test/nn/test_qlinear.py::test_move_qlinear

I need to check why these tests are not passing anymore with this PR :

FAILED test/tensor/ops/test_quantized_dispatch.py::test_to_device[cuda] - AssertionError: assert 'cpu' == 'cuda'
FAILED test/tensor/ops/test_quantized_dispatch.py::test_softmax[cuda-5-5-1] - AssertionError
FAILED test/tensor/ops/test_quantized_dispatch.py::test_softmax[cuda-5-5-10] - AssertionError
FAILED test/tensor/ops/test_quantized_dispatch.py::test_softmax[cuda-32-32-1] - AssertionError
FAILED test/tensor/ops/test_quantized_dispatch.py::test_softmax[cuda-32-32-10] - AssertionError
FAILED test/tensor/ops/test_quantized_dispatch.py::test_softmax[cuda-10-32-1] - AssertionError
FAILED test/tensor/ops/test_quantized_dispatch.py::test_softmax[cuda-10-32-10] - AssertionError

cc @dacorvo

[SunMarc]

The following works even before this PR. This is why you never had this specific issue.

[dacorvo]

There is already a test for this here: https://github.com/huggingface/quanto/blob/6302171b7569a3fd86f31e1731d01d390d1eb557/test/tensor/ops/test_quantized_dispatch.py#L8

[SunMarc]

I wanted to return super().to(*args, **kwargs) but it was causing weird behavior with tests using QBitsTensor and it was calling __torch_function__ after . To reproduce, return super().to(*args, **kwargs) and run python -m pytest -sv test/nn/test_qlinear.py::test_move_qlinear

[dacorvo]

Tensor subclasses are a very special beasts: you should not override the base Tensor methods that way, and instead do it through the dispatch.

[dacorvo]

I am not sure there is actually an issue with QTensor, as .to() is already correctly dispatched here:

https://github.com/huggingface/quanto/blob/6302171b7569a3fd86f31e1731d01d390d1eb557/quanto/tensor/ops.py#L74

It won't work for QBitsTensor that has an extra zeropoint inner tensor, but it should be fixed at the dispatch stage, and not directly in the class declaration IMHO.

[dacorvo]

I can reproduce the issue when moving a module. This happens because the move happens in two steps:

• move the QTensor (this calls the dispatch): t -> new_t
• assign it back to the module param.

The second step does not replace the original tensor, though: instead it does a shallow copy by doing t.data = new_t.

This results in a weird situation indeed where the moved tensor is stored as an attribute of the original one ...

[dacorvo]

Closing as obsoleted by #90