refactor common used toy model #2729

namgyu-youn · 2025-08-11T08:49:18Z

Summary:
Integrate commonly used single/multi-linear toy models and refactor them across the codebase (src/test/benchmark/tutorial).

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

Test Plan: CI

Integrates common used toy model and refactor across TorchAO (ao/test/benchmark/tutorial) - fix: pytorch#2078 Test Plan: CI

pytorch-bot · 2025-08-11T08:49:21Z

🔗 Helpful Links

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

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

jainapurva · 2025-08-11T16:24:52Z

@namgyu-youn thanks for taking up this effort

jainapurva · 2025-08-11T16:35:52Z

torchao/testing/model_architectures.py

-        self.linear1 = torch.nn.Linear(k, n, bias=False).to(dtype)
+        self.linear1 = torch.nn.Linear(m, n, bias=False)
+        self.linear2 = torch.nn.Linear(n, k, bias=False)
+        self.linear3 = torch.nn.Linear(k, 1, bias=False)


Please create a separate model for two linear layers. This model for single linear layer is used in benchmarking run on CI.

@jainapurva I prefer to define ToySingleLinearModel and ToyMultiLinearModel for a future update as you mentioned, but how about reverting benchmark_aq.py?

Unit tests (e.g., test_quant_api.py, test_awq.py) are using single/multiple layers in a mixed manner, and using only multiple layers might be an update. If this makes sense, benchmark_aq.py would be the only case using single layers. Let me know which one aligns better.

ToySingleLinearModel and ToyMultiLinearModel sounds good. Please ensure all the tests are running smoothly for it.
For benchmark_aq.py you can add the bias parameter as the last arg in init and set it to False by default. In addition to this, ToySingleLinearModel is used in running .github/workflows/run_microbenchmarks.yml. It uses the create_model_and_input_data, please ensure that method is running smoothly, and is updated as per the new toy models.

Sorry for opening the PR without checking it. And I will move into your suggestion; thanks for your leading.

jainapurva · 2025-08-11T22:25:54Z

@namgyu-youn Please feel free to divide this into multiple PRs if it's too many changes.

Integrate commonly used single/multi-linear toy models and refactor them across the codebase (src/test/benchmark/tutorial). - fix: pytorch#2078 Test Plan: CI

jainapurva · 2025-08-14T17:21:57Z

@namgyu-youn There are some merge conflicts in the branch. Please rebase it onto main. If needed, I can help with that.

namgyu-youn · 2025-08-17T05:53:56Z

@jainapurva Could you take a look at this PR? It passed CI after resolving the merge conflict.

jerryzh168 · 2025-08-18T22:57:17Z

torchao/testing/model_architectures.py

+        self.linear3 = torch.nn.Linear(k, 64, bias=has_bias)
+
+    def example_inputs(
+        self, batch_size=1, sequence_length=10, dtype=torch.float32, device="cpu"


nit: should we move dtype and device to __init__ as well to be consistent?

If there is a plan to expand the toy model like backward, we can consider moving them to __init__. But I am fine to keep this because it has slightly more brevity, and there is no plan to expand it.

what I meant is that, linear in init should have dtype and device as well, and it doesn't make sense to define linear modules in one device/dtype but get example_inputs from another dtype/device? so might be easier just to define these in init and not worry about them in example_inputs?

Oh I missed it. Updating init is much better, thanks.

jerryzh168 · 2025-08-21T23:05:36Z

torchao/testing/model_architectures.py

+
+
+class ToyMultiLinearModel(torch.nn.Module):
+    def __init__(self, m=512, n=256, k=128, has_bias=True):


I feel m, n, k should be required actually (also probably change m, n, k naming a bit, since it's easily confused with the shapes of linear itself)

also do we need 3 linears? can this be 2 linears and rename to TwoLinearModel to make it clearer?

That means the user should input m, n, k whenever ToyMultiLinearModel is called, right? In my thought, (m, n, k) can be renamed to (input_dim, hidden_dim, output_dim). Let me know if there is better one.

Also in the old version, there were two (test_awq.py and test_smoothquant.py) scripts related to performance (error range; AWQ and SmoothQuant). But since they are quiet far away from the real benchmark, I am fine to go with 2-linears for more brevity. ToyTwoLinearModel sounds good to me.

jerryzh168 · 2025-08-26T01:07:47Z

docs/source/serialization.rst

-          x = self.linear1(x)
-          x = self.linear2(x)
-          return x
+  from torchao.testing.model_architectures import ToyTwoLinearModel


can you revert the changes for this? I think it's better to have this tutorial self contained

Yes keeping them for tutorial sounds good to me, I will revert it.

jerryzh168 · 2025-08-26T01:08:02Z

docs/source/quick_start.rst

@@ -29,19 +29,9 @@ First, let's set up our toy model:

  import copy
  import torch
+  from torchao.testing.model_architectures import ToyTwoLinearModel


jerryzh168 · 2025-08-26T01:08:38Z

torchao/testing/model_architectures.py

+    """Single linear for m * k * n problem size"""
+
+    def __init__(
+        self, m=64, n=32, k=64, has_bias=False, dtype=torch.float, device="cuda"


default dtype should probably be torch.bfloat16 I feel

jerryzh168 · 2025-08-26T01:08:56Z

torchao/testing/model_architectures.py

+
+
+class ToyTwoLinearModel(torch.nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, has_bias=False):


dtype and device?

oh I missed it, thanks.

jerryzh168 · 2025-08-26T01:09:24Z

torchao/testing/model_architectures.py

@@ -179,7 +220,7 @@ def create_model_and_input_data(
        m, k, n (int): dimensions of the model and input data
    """
    if model_type == "linear":
-        model = ToyLinearModel(k, n, high_precision_dtype).to(device)
+        model = ToyTwoLinearModel(k, n, high_precision_dtype).to(device)


arg seems to be wrong here?

I misunderstood its workflow. This if-else and test_model_architecture.py should be updated using the following:

model, input_data = create_model_and_input_data( "linear", 10, 64, 32, device=device )

jerryzh168 · 2025-08-26T01:09:49Z

tutorials/calibration_flow/static_quant.py

@@ -284,7 +265,7 @@ def test_static_quant(target_dtype: torch.dtype, mapping_type: MappingType):
    weight_obs = AffineQuantizedMinMaxObserver(
        mapping_type,
        target_dtype,
-        granularity=PerAxis(axis=0),
+        granularity=PerTensor(),


why is this changed?

It was my misundertstooding while fixing obersver's input shape. I will revert it with related workflows.

jerryzh168 · 2025-08-26T01:10:05Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

@@ -113,7 +102,7 @@ def test_fp8_linear_variants(
            input_tensor = torch.randn(*M, K, dtype=dtype, device="cuda")

            # Create a linear layer with bfloat16 dtype
-            model = ToyLinearModel(K, N).eval().to(dtype).to("cuda")
+            model = ToyTwoLinearModel(K, 64, N).eval().to(dtype).to("cuda")


Unlike ToyLinearModel (old), ToyTwoLinearModel uses input_dim (K), hidden_dim (64), output_dim (N); following is the same case.

jerryzh168 · 2025-08-26T01:10:24Z

test/quantization/quantize_/workflows/float8/test_float8_tensor.py

@@ -222,7 +211,7 @@ def test_kernel_preference_numerical_equivalence(self, granularity, sizes):
        dtype = torch.bfloat16
        input_tensor = torch.randn(*M, K, dtype=dtype, device="cuda")
        # Create a linear layer with bfloat16 dtype
-        model = ToyLinearModel(K, N).eval().to(dtype).to("cuda")
+        model = ToyTwoLinearModel(K, 64, N).eval().to(dtype).to("cuda")


jerryzh168 · 2025-08-26T01:11:10Z

please run the changed tests locally as well

Summary:

b2e7f54

Integrates common used toy model and refactor across TorchAO (ao/test/benchmark/tutorial) - fix: pytorch#2078 Test Plan: CI

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 11, 2025

jainapurva self-requested a review August 11, 2025 16:24

jainapurva added the topic: not user facing Use this tag if you don't want this PR to show up in release notes label Aug 11, 2025

fix ruff

c5faa07

jainapurva reviewed Aug 11, 2025

View reviewed changes

namgyu-youn marked this pull request as draft August 11, 2025 22:23

separate single/multi linear toy model

ddeb027

namgyu-youn marked this pull request as ready for review August 12, 2025 01:04

namgyu-youn requested a review from jainapurva August 12, 2025 01:04

Summary:

2aafd64

Integrate commonly used single/multi-linear toy models and refactor them across the codebase (src/test/benchmark/tutorial). - fix: pytorch#2078 Test Plan: CI

Merge branch 'main' into refactor-toymodel

68e4482

namgyu-youn marked this pull request as draft August 16, 2025 16:08

fix CI error after rebase

6e88012

namgyu-youn marked this pull request as ready for review August 17, 2025 05:52

jerryzh168 reviewed Aug 18, 2025

View reviewed changes

jerryzh168 reviewed Aug 21, 2025

View reviewed changes

update 3-linear model to 2-linear model

6fd9672

namgyu-youn requested a review from jerryzh168 August 22, 2025 10:27

Merge branch 'main' into refactor-toymodel

98dd997

jerryzh168 reviewed Aug 26, 2025

View reviewed changes

namgyu-youn added 4 commits August 26, 2025 16:10

revert: observer shape

1656126

revert: toy model for tutorials

994b507

update dtype, device handling in ToyTwoLinearModel

0ced363

fix: test module for create_model_and_input_data()

6b03dc3

namgyu-youn requested a review from jerryzh168 August 26, 2025 15:42



		class ToyMultiLinearModel(torch.nn.Module):
		def __init__(self, m=512, n=256, k=128, has_bias=True):



		class ToyTwoLinearModel(torch.nn.Module):
		def __init__(self, input_dim, hidden_dim, output_dim, has_bias=False):

refactor common used toy model #2729

Are you sure you want to change the base?

refactor common used toy model #2729

Uh oh!

Conversation

namgyu-youn commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

Uh oh!

jainapurva commented Aug 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

namgyu-youn Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jainapurva commented Aug 11, 2025

Uh oh!

jainapurva commented Aug 14, 2025

Uh oh!

namgyu-youn commented Aug 17, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jerryzh168 Aug 21, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

namgyu-youn Aug 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jerryzh168 commented Aug 26, 2025

Uh oh!

Uh oh!

namgyu-youn commented Aug 11, 2025 •

edited

Loading

pytorch-bot bot commented Aug 11, 2025 •

edited

Loading

namgyu-youn Aug 11, 2025 •

edited

Loading

jerryzh168 Aug 21, 2025 •

edited

Loading

namgyu-youn Aug 22, 2025 •

edited

Loading