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). Integrated toy model (ToyTwoLinearModel) no longer uses default device, dtype.

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

Test Plan:

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
test/sparsity/test_fast_sparse_training.py

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


  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

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

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

thanks please see a few final comments inline

jerryzh168

looks good, thanks for addressing all the comments!

jerryzh168 · 2025-10-08T01:39:50Z

seems there are some new code landed in test_quant_api.py please fix these as well, see https://github.com/pytorch/ao/actions/runs/18330669725/job/52205257444?pr=2729

namgyu-youn · 2025-10-08T01:59:29Z

seems there are some new code landed in test_quant_api.py please fix these as well, see https://github.com/pytorch/ao/actions/runs/18330669725/job/52205257444?pr=2729

Could you re-run the workflow? There's no merge conflict on my end; it might be related to CI cache issue.

jerryzh168 · 2025-10-08T02:12:31Z

not merge conflict, but new code that's added that needs refactor

namgyu-youn · 2025-10-08T18:35:35Z

not merge conflict, but new code that's added that needs refactor

done, please check this PR

jerryzh168 · 2025-10-08T19:34:20Z

still some issue with formatting, please try pre-commit install and pre-commit run #1448 (comment)

jerryzh168

formatting

namgyu-youn · 2025-10-20T07:01:39Z

@jerryzh168 please take a look into this PR; I hope this PR is ready to land.

jerryzh168 · 2025-10-30T22:44:48Z

looks like there are some failures and merge conflict, can you rebase

namgyu-youn · 2025-10-31T03:08:54Z

looks like there are some failures and merge conflict, can you rebase

Fixed rebase, but CI failure seems unrelated to this PR, I think

jerryzh168 · 2025-10-31T17:31:03Z

some test errors, can you take a look?

after the error is fixed, please also rebase

namgyu-youn · 2025-10-31T19:51:22Z

@jerryzh168 I'll restart this work with smaller chunks and ping again because this PR is chained across codebase, which made it hard for me to fix CI. This should be better for rebase issue I believe.

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

jerryzh168 requested changes Oct 7, 2025

View reviewed changes

update toy model args

ddb6b0d

jerryzh168 approved these changes Oct 8, 2025

View reviewed changes

namgyu-youn requested a review from jerryzh168 October 8, 2025 01:19

namgyu-youn added 2 commits October 9, 2025 03:32

add new code with toy model

ef6c3f4

Merge branch 'main' into refactor-toymodel

7630988

jerryzh168 approved these changes Oct 8, 2025

View reviewed changes

jerryzh168 requested changes Oct 8, 2025

View reviewed changes

fix pre-commit format

d89f1d6

namgyu-youn requested a review from jerryzh168 October 9, 2025 03:47

namgyu-youn added 2 commits October 17, 2025 01:20

Merge branch 'main' into refactor-toymodel

0c1d4d7

fix after rebase

86a81a9

namgyu-youn marked this pull request as draft October 20, 2025 06:32

fix lint

22d341b

namgyu-youn marked this pull request as ready for review October 20, 2025 06:59

namgyu-youn added 2 commits October 31, 2025 09:09

Merge branch 'main' into refactor-toymodel

c6cc73c

fix after rebase

9110fd1

Merge branch 'main' into refactor-toymodel

0bb1938

namgyu-youn closed this Oct 31, 2025

namgyu-youn mentioned this pull request Nov 1, 2025

Update common used toy linear model #3275

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

refactor common used toy model #2729

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!

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

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!

jerryzh168 left a comment

Choose a reason for hiding this comment

Uh oh!

jerryzh168 left a comment

Choose a reason for hiding this comment

Uh oh!

jerryzh168 commented Oct 8, 2025

Uh oh!

namgyu-youn commented Oct 8, 2025

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 29, 2025 •

edited

Loading

namgyu-youn Aug 26, 2025 •

edited

Loading

namgyu-youn commented Oct 31, 2025 •

edited

Loading

namgyu-youn commented Oct 31, 2025 •

edited

Loading