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.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

ROCm MI2xx CI/CD workflows failing due to : download from https://api.github.com/repos/pytorch/pytorch timed out.

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

this is still not updated

Oh I missed it, thanks for the reminder.

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 )

Therefore, we should not use toy model here.

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.

for old model, it's in_features = K, out_features = N
mapping to ToyTwoLInearModel should be input_dim = K, hidden_dim = N, output_dim = K, right?

Yes this is right, thanks for correcting these. I will fix all these 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")


Yes all dimensions in ToyTwoLinearModel should be updated, thanks.

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

please run the changed tests locally as well

jerryzh168 · 2025-08-28T17:17:35Z

docs/source/serialization.rst

@@ -46,7 +46,7 @@ Here is the serialization and deserialization flow::
      state_dict = torch.load(f)

  with torch.device("meta"):
-      m_loaded = ToyLinearModel(1024, 1024, 1024).eval().to(dtype)
+      m_loaded = ToyTwoLinearModel(1024, 1024, 1024).eval().to(dtype)


this has to be reverted as well?

Yes it was reverted at 994b507.

Sorry I misunderstood it. What you mean is revert its name also, right? We can keep ToyLinearModel for the tutorials.

yeah, otherwise the tutorial code won't run

jerryzh168 · 2025-08-28T17:18:26Z

scripts/quick_start.py

-
-
-model = ToyLinearModel(1024, 1024, 1024).eval().to(torch.bfloat16).to("cuda")
+model = ToyTwoLinearModel(1024, 1024, 1024).eval().to(torch.bfloat16).to("cuda")


this should be reverted as well I think, since it seems like to be a copy of the quick_start.rst

jerryzh168 · 2025-08-28T17:20:33Z

torchao/testing/model_architectures.py

+        output_dim,
+        has_bias=False,
+        dtype=torch.bfloat16,
+        device="cpu",


default device should also be cuda here I think

jerryzh168 · 2025-08-28T17:21:55Z

torchao/testing/model_architectures.py

+        if device is None:
+            device = self.device
+
+        if sequence_length is not None:


is this used? this can be done in callsite as well right?

Yes reverting sequence_length and removing this conditioner is much better.

sorry I missed AWQ. It is used at test/prototype/test_awq.py, so we have to keep sequence_length for AWQ.

jerryzh168 · 2025-08-28T17:22:43Z

torchao/testing/model_architectures.py

+        )
+
+    def example_inputs(
+        self, batch_size=1, sequence_length=None, dtype=None, device=None


also not sure if we need dtype and device arg here

Yes we only need init for ags

jerryzh168 · 2025-08-28T17:23:27Z

test/dtypes/test_affine_quantized_float.py

@@ -122,7 +109,7 @@ def test_fp8_linear_variants(
            }

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


original seems to be (K, N, K)? according to L54-55 in original file

same for the many of the changes in this file, I think we can match the original

Yes it is (K, N, K), not (M, N, K); thanks for correcting these.

jerryzh168 · 2025-08-28T18:27:11Z

torchao/testing/model_architectures.py

-# TODO: Refactor torchao and tests to use these models
-class ToyLinearModel(torch.nn.Module):
-    def __init__(self, k=64, n=32, dtype=torch.bfloat16):
+class ToySingleLinearModel(torch.nn.Module):


I also feel it's important for ToySingleLinearModel and ToyTwoLinearModel to have consistent APIs, right now they are not:

batch_size is passed in example_input() function for ToyTwoLinearModel but not ToySingleLinearModel

dtype, device arg list mismatch for init and example_inputs

default values for dtype, device should also match

what example_inputs returns also mismatch

Sound good to me. But because single-linear model is chained with benchmarks/benchmark_aq.py, it would be a little bit harder. ~~In other words, updating dtype, device, or whatever in single-linear model might break CI, require to update docs, etc.~~ Lcukily, ToySingleLinearModel used same dtype, device as ToyTwoLinearModel, so we can do it without considering CI.

Personally, defining ToyLinearModel(layer=1,2, ..., n) seems more general. How about moving to ToyLinearModel(layer=n) after this PR? But this is low-priority I guess.

we can extend to layer n in the future if needed, keep 2 for now for simplicity would be better I think

namgyu-youn · 2025-08-29T09:57:37Z

test/prototype/test_awq.py

@@ -108,7 +85,7 @@ def test_awq_functionality(self):

        loss_awq = (ref_out - awq_out).pow(2).mean().item()
        loss_base = (ref_out - baseline_out).pow(2).mean().item()
-        assert loss_awq < loss_base
+        assert loss_awq < loss_base * 1.1


Since this is edge case (toy model architecture is quiet different), error range is adjusted for passing CI. We can try only checking loss_awq is generated (no matter error range), as discussed in #2728 (comment) for more brevity

I'm not sure we can do that, even the model changed, the loss should still be smaller I think, since that's waht awq is optimizing for

Maybe higher error cames from lower (3->2) layers. Because AWQ uses weight distribution in this implementation, 2-layers might not adequate to compute distribution, making AWQ hard to learn.

namgyu-youn · 2025-08-29T11:00:10Z

torchao/testing/model_architectures.py

+        if device is None:
+            device = self.device
+
+        if sequence_length is not None:


sorry I missed AWQ. It is used at test/prototype/test_awq.py, so we have to keep sequence_length for AWQ.

namgyu-youn · 2025-08-29T11:36:56Z

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")


Yes all dimensions in ToyTwoLinearModel should be updated, thanks.

namgyu-youn · 2025-08-29T12:15:36Z

torchao/testing/model_architectures.py

+        )
+
+    def example_inputs(
+        self, batch_size=1, sequence_length=None, dtype=None, device=None


Yes we only need init for ags

jerryzh168 · 2025-08-30T02:44:25Z

torchao/testing/model_architectures.py

+        if sequence_length is not None:
+            return [
+                torch.randn(
+                    1, self.linear1.in_features, dtype=self.dtype, device=self.device


this is different from the previous code in AWQ I think, sequence_length is not used at all here

I feel we can just copy paste this to awq, instead of complicating the implementation of ToyTwoLinearModel here

Ok I am fine to revert it for this edge-case.

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

jerryzh168 reviewed Aug 28, 2025

View reviewed changes

namgyu-youn added 2 commits August 29, 2025 20:59

revert: toy model for tutorials

6b4eaa8

fix: uniform args (device & dtype) in ToyModel

ee7b0f4

namgyu-youn commented Aug 29, 2025

View reviewed changes

namgyu-youn requested a review from jerryzh168 August 29, 2025 13:37

jerryzh168 reviewed Aug 30, 2025

View reviewed changes

namgyu-youn added 2 commits August 30, 2025 14:42

remove overused args: sequence_length

c8320a7

revert edge-case to source: test_awq.py

b6a752e

namgyu-youn requested a review from jerryzh168 August 30, 2025 09:17

namgyu-youn mentioned this pull request Aug 30, 2025

Make SmoothQuant more General #2728

Open



		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):



		model = ToyLinearModel(1024, 1024, 1024).eval().to(torch.bfloat16).to("cuda")
		model = ToyTwoLinearModel(1024, 1024, 1024).eval().to(torch.bfloat16).to("cuda")

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

❗ 1 Active SEVs

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!

namgyu-youn Aug 26, 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!

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

namgyu-youn Aug 26, 2025 •

edited

Loading

namgyu-youn Aug 29, 2025 •

edited

Loading

namgyu-youn Aug 29, 2025 •

edited

Loading