avoid .to() followed by inplace mutation to appease export

[bdhirsh]

Export currently has a restriction, where inplace mutations on the output of a call to aten.to() results in an error.

We should really lift this restriction, but in the meantime, this PR avoids hitting that problem in the quantization logic.

I tested by adding the following lines to the bottom of https://github.com/hustvl/ViTMatte/blob/main/modeling%2Fbackbone%2Fvit.py and confirming that export did not error:

quantize_(m, int8_dynamic_activation_int8_weight())
m = unwrap_tensor_subclass(m)

x = torch.randn(4, 3, 128, 128)
m = torch.export.export(m, (x,))
m2 = m.run_decompositions()
print(m2)

Stack from ghstack (oldest at bottom):

• -> avoid .to() followed by inplace mutation to appease export #1387

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 08bbd50 with merge base 8a805d0 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
  test/integration/test_integration.py::TestSubclass::test_int8_dynamic_quant_subclass_api_5_cuda

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

[bdhirsh]

fyi @tugsbayasgalan @yushangdi

[bhack]

/opt/conda/lib/python3.11/site-packages/torchao/dtypes/uintx/block_sparse_layout.py:        y += bias
/opt/conda/lib/python3.11/site-packages/torchao/dtypes/uintx/plain_layout.py:        y += bias
/opt/conda/lib/python3.11/site-packages/torchao/dtypes/uintx/semi_sparse_layout.py:        y += bias
/opt/conda/lib/python3.11/site-packages/torchao/dtypes/uintx/tensor_core_tiled_layout.py:        y += bias
/opt/conda/lib/python3.11/site-packages/torchao/dtypes/uintx/uint4_layout.py:                y += bias
/opt/conda/lib/python3.11/site-packages/torchao/prototype/quantization/autoquant_v2.py:            y += bias
/opt/conda/lib/python3.11/site-packages/torchao/prototype/quantization/autoquant_v2.py:            y += bias
/opt/conda/lib/python3.11/site-packages/torchao/quantization/autoquant.py:            y += bias
/opt/conda/lib/python3.11/site-packages/torchao/quantization/autoquant.py:            y += bias
/opt/conda/lib/python3.11/site-packages/torchao/quantization/subclass.py:            y += bias
/opt/conda/lib/python3.11/site-packages/torchao/quantization/subclass.py:            y += bias
/opt/conda/lib/python3.11/site-packages/torchao/quantization/weight_only.py:            y += self.bias
```

[jerryzh168]

thanks, we can't use ghstack in torchao btw, you can use https://github.com/modularml/stack-pr or just normal git push

[jerryzh168]

I'm not sure if changing this would have perf implications actually, probably need to benchmark a bit for the popular code paths

[bhack]

I think it is better to have export coverage in the CI if we can.

[bhack]

Related pytorch/pytorch#138606

[bhack]

/cc @yushangdi

[yushangdi]

Related pytorch/pytorch#138606

fyi, @jerryzh168 @bhack , the approach in the PR linked above doesn't work. I'm trying out a new approach, mostly this:

if we have a.to(b), and b has a different dtype with a, then it must be a copy. In this case, we do not need to freeze the tensor. Instead, we use torch.ops.aten._assert_tensor_metadata.default to ensure that a must not have the same dtype as b.

cc @tugsbayasgalan