Skip to content

Metal backend: fail test_modules.py if torchao can't be imported #17236

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
manuelcandales wants to merge 5 commits into
base: gh/manuelcandales/153/head
Choose a base branch
from
+5 −20
Open

Metal backend: fail test_modules.py if torchao can't be imported #17236

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5774967
Update
manuelcandales Feb 5, 2026
a103c16
Update
manuelcandales Feb 5, 2026
58384d7
Update
manuelcandales Feb 5, 2026
7f258b0
Update
manuelcandales Feb 5, 2026
524580b
Update
manuelcandales Feb 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 5 additions & 20 deletions backends/apple/metal/tests/test_modules.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,22 +24,17 @@

import numpy as np
import torch

# Need to import to load the torchao metal ops
import torchao.experimental.ops.mps # noqa: F401

from executorch.backends.apple.metal.metal_backend import MetalBackend
from executorch.backends.apple.metal.metal_partitioner import MetalPartitioner
from executorch.exir import to_edge_transform_and_lower
from torch import nn
from torch.export import export
from torch.nn.attention import SDPBackend

try:
# Need to import to load the ops
import torchao.experimental.ops.mps # noqa: F401
from torchao.experimental.quant_api import UIntxWeightOnlyConfig
from torchao.quantization.quant_api import quantize_

TORCHAO_AVAILABLE = True
except ImportError:
TORCHAO_AVAILABLE = False
from torchao.experimental.quant_api import quantize_, UIntxWeightOnlyConfig
Copy link
Contributor

@mergennachin mergennachin Feb 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

from torchao.quantization.quant_api import quantize_



# Check if MPS is available for export tests
Expand Down Expand Up @@ -245,7 +240,6 @@ def forward(self, x: torch.Tensor):
"rtol_float32": 5e-2,
"atol_bfloat16": 1e-1,
"rtol_bfloat16": 1e-1,
"skip": not TORCHAO_AVAILABLE,
}


Expand All @@ -270,7 +264,6 @@ def forward(self, x: torch.Tensor):
"rtol_float32": 5e-2,
"atol_bfloat16": 1e-1,
"rtol_bfloat16": 1e-1,
"skip": not TORCHAO_AVAILABLE,
}


Expand All @@ -295,7 +288,6 @@ def forward(self, x: torch.Tensor):
"rtol_float32": 5e-2,
"atol_bfloat16": 1e-1,
"rtol_bfloat16": 1e-1,
"skip": not TORCHAO_AVAILABLE,
}


Expand All @@ -320,7 +312,6 @@ def forward(self, x: torch.Tensor):
"rtol_float32": 5e-2,
"atol_bfloat16": 1e-1,
"rtol_bfloat16": 1e-1,
"skip": not TORCHAO_AVAILABLE,
}


Expand All @@ -345,7 +336,6 @@ def forward(self, x: torch.Tensor):
"rtol_float32": 5e-2,
"atol_bfloat16": 1e-1,
"rtol_bfloat16": 1e-1,
"skip": not TORCHAO_AVAILABLE,
}


Expand Down Expand Up @@ -697,11 +687,6 @@ def quantize_model(model: nn.Module, qlinear: str, qlinear_group_size: int = 32)
- "fpa4w": Floating point activation, 4-bit weight (Metal backend)
qlinear_group_size: Group size for quantization (default: 32).
"""
if not TORCHAO_AVAILABLE:
raise RuntimeError(
"torchao is not available. Install torchao to use quantization."
)

if qlinear == "fpa4w":
linear_config = UIntxWeightOnlyConfig(
group_size=qlinear_group_size,
Expand Down
Loading