update toy model to use global variants

[namgyu-youn]

Summary:
This PR is smaller chunks of #2729. Instead of implementing a linear model in each test code, we want to support its API for developers' utility and consistency. The updated toy model requires device and dtype to help clearly identify.

• related to: Refactor torchao and tests to use model architectures from torchao.testing.model_architectures #2078

Test Plan:
benchmarks/benchmark_aq.py
test/dtypes/test_affine_quantized_float.py

Future Plan:
For a full update, we also have to update the following tests, the same as this PR:

• test/integration/test_integration.py
• test/quantization/quantize_/workflows/float8/test_float8_tensor.py
• test/quantization/test_quant_api.py

Co-authored-by: Jerry Zhang jerryzh168@gmail.com

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[namgyu-youn]

cc @jerryzh168 :)

[namgyu-youn]

@pytorchbot label "topic: not user facing" "topic: for developers"

[jerryzh168]

I think we can prioritize [test/quantization/quantize_/workflows/float8/test_float8_tensor.py](https://github.com/pytorch/ao/blob/main/test/quantization/quantize_/workflows/float8/test_float8_tensor.py) others might have to be deprecated

[namgyu-youn]

I think we can prioritize [test/quantization/quantize_/workflows/float8/test_float8_tensor.py](https://github.com/pytorch/ao/blob/main/test/quantization/quantize_/workflows/float8/test_float8_tensor.py) others might have to be deprecated

Actually the order is intentional for the small-to-large process. So the order will be:

1. Base model: torchao/testing/model_architecture.py
2. Old api (we are here)
3. New tensor subclass: Int8Tensor, Float8Tensor
4. Quantization api integration: test/quantization/test_quant_api.py
5. Full integration: test/integration/test_integration.py

[jerryzh168]

OK, please check the CI errors

[namgyu-youn]

@pytorchbot label "topic: not user facing" "topic: for developers"

[namgyu-youn]

OK, please check the CI errors

Done, please take a look at this PR.