update toy model to use api

[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, same as this PR:

• test/dtypes/test_affine_quantized_float.py
• 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 [email protected]

[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.

❌ 3 New Failures, 1 Unrelated Failure

As of commit b61a7b0 with merge base 1a9b6f4 ():

NEW FAILURES - The following jobs have failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.
• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh)
  test/dtypes/test_affine_quantized_float.py::TestAffineQuantizedFloat8Compile::test_serialization_mode_weight-only
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh)
  test/dtypes/test_affine_quantized_float.py::TestAffineQuantizedFloat8Compile::test_serialization_mode_weight-only

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 (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh) (trunk failure)
  test/test_low_bit_optim.py::TestOptim::test_param_groups_optim_name_AdamFp8_device_cpu

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

Oh actually the order is intentional for 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