Metal backend: enable quantization in test_modules

[manuelcandales]

This pull request adds support for quantized model testing (specifically int4 quantization) to the Metal backend module tests, using torchao. It introduces the ability to specify quantization options in the module registry, applies quantization to linear layers as needed, and updates the testing pipeline to compare quantized models against unquantized references when appropriate. A new test model with int4 quantization is also added.

Quantization Support and Configuration:

• Added imports and availability checks for torchao and its quantization APIs, enabling conditional quantization support in tests.
• Expanded the module registry documentation to describe new quantization-related options (qlinear, qlinear_group_size, compare_to_unquantized) and how to use them for int4 quantization.

Test Model and Registry Enhancements:

• Introduced LinearNoBiasInt4, a linear layer test model without bias and with int4 quantization, and registered it in MODULE_REGISTRY with appropriate quantization and tolerance settings.

Quantization Application Logic:

• Enhanced get_model_and_inputs to accept quantization arguments, apply quantization to models as specified in the registry, and added the quantize_model helper to perform quantization using torchao. [1] [2]

Testing Pipeline Updates:

• Updated export_model_to_files to optionally compare quantized model outputs to unquantized references, handling device placement and ensuring correct output computation for quantized models.
• Modified the main test routine to determine whether to compare quantized models to unquantized references based on the registry configuration.

[manuelcandales]

Stack from ghstack (oldest at bottom):

• Metal backend: fix bugs in qmv_impl #17188
• Metal backend: import qmv_impl from MLX #17187
• Metal backend: int4 refactor checks and logs #17186
• Metal backend: test 4-bit linear with bias #17127
• -> Metal backend: enable quantization in test_modules #17118
• Metal backend: support 4-bit linear #17117
• Metal backend: remove parakeet linear decomp #17116
• Metal backend: enable linear with bias #17115
• Metal backend: test modules #17076

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/17118

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 2 Unrelated Failures

As of commit fec15bc with merge base ba6de95 ():

NEW FAILURES - The following jobs have failed:

• pull / unittest-arm-backend-with-no-deps (test_pytest_models_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t a6e00651f7610d05e8c877c6949d59520a43b38a2a5e18ea50de73e77be53974 /exec failed with exit code 1
• pull / unittest-buck / linux / linux-job (gh)
  RuntimeError: Command docker exec -t cbc54fc0accd0313384c973a92ed7b8f938a41ec7f756c3d6130343a1c80e533 /exec failed with exit code 3

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-custom-ops-linux / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 1

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-buck / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 3

This comment was automatically generated by Dr. CI and updates every 15 minutes.