Arm backend: Add QAT support to test framework and model tests (#22741)

### Summary

Adds an `is_qat` option to the shared TOSA and VGF test pipelines so
quantization-aware training follows the existing export and lowering
flow while PTQ remains the default. Adds QAT smoke coverage for the
current FP32 NSS and NFRU models through TOSA and VGF.


### Test plan

Gerrit patch set 13 passed CI (`Verified +1`) and received `Code-Review
+2` and `IP-Review +1`. Also checked the patch with `git diff --check`.

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Author: usamahz

backends/arm/MODELS.md

@@ -20,7 +20,8 @@
 - Some popular torch.nn.modules models (NN modules)
 - Some popular torch ops (Torch Functions)
 - T5 (T5 for conditional generation)
-- Neural Super Sampler (NSS)
+- Neural Frame Rate Upscaling (NFRU), including QAT smoke coverage
+- Neural Super Sampling (NSS), including QAT smoke coverage
 - Phi-3
 - ResNet 18
 - ResNet-8

backends/arm/test/models/test_nfru.py

@@ -0,0 +1,92 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+
+os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
+
+import pytest
+import torch
+
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    TosaPipelineINT,
+    VgfPipeline,
+)
+from huggingface_hub import hf_hub_download
+from ng_model_gym.usecases.nfru.model.nfru_v1_nn import (  # type: ignore[import-not-found,import-untyped]
+    NFRUAutoEncoder,
+)
+
+
+_RELEASE_REFS = (
+    os.environ.get("GITHUB_REF", ""),
+    os.environ.get("GITHUB_REF_NAME", ""),
+    os.environ.get("GITHUB_BASE_REF", ""),
+)
+_IS_FROZEN_RELEASE = any(
+    ref.removeprefix("refs/heads/").startswith("release/") for ref in _RELEASE_REFS
+)
+pytestmark = pytest.mark.skipif(
+    _IS_FROZEN_RELEASE,
+    reason="NFRU tests depend on resources fetched from main.",
+)
+
+_NFRU_QAT_INPUTS = (torch.ones((1, 16, 64, 64)),)
+
+
+def nfru() -> NFRUAutoEncoder:
+    weights = hf_hub_download(  # nosec B615
+        repo_id="Arm/neural-frame-rate-upscaling",
+        filename="nfru_v1_fp32.pt",
+        revision="main",
+        cache_dir=os.environ.get("RUNNER_TEMP"),
+    )
+    checkpoint = torch.load(
+        weights,
+        map_location=torch.device("cpu"),
+        weights_only=True,
+    )["model_state_dict"]
+    prefix = "network.auto_encoder."
+    assert all(key.startswith(prefix) for key in checkpoint)
+
+    model = NFRUAutoEncoder()
+    model.load_state_dict(
+        {key.removeprefix(prefix): value for key, value in checkpoint.items()},
+        strict=True,
+    )
+    return model
+
+
+def test_nfru_qat_tosa_INT() -> None:
+    pipeline = TosaPipelineINT(
+        nfru(),
+        _NFRU_QAT_INPUTS,
+        aten_op=[],
+        exir_op=[],
+        per_channel_quantization=False,
+        use_to_edge_transform_and_lower=True,
+        is_qat=True,
+        frobenius_threshold=None,
+        cosine_threshold=None,
+    )
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+def test_nfru_qat_vgf_INT() -> None:
+    pipeline = VgfPipeline(
+        nfru(),
+        _NFRU_QAT_INPUTS,
+        aten_op=[],
+        exir_op=[],
+        run_on_vulkan_runtime=True,
+        quantize=True,
+        per_channel_quantization=False,
+        use_to_edge_transform_and_lower=True,
+        is_qat=True,
+        tosa_version="TOSA-1.0+INT",
+    )
+    pipeline.run()

backends/arm/test/models/test_nss.py

@@ -212,3 +212,38 @@ def test_nss_vgf_INT(use_real_data):
     if use_real_data:
         _set_nss_calibration_samples(pipeline)
     pipeline.run()
+
+
+def test_nss_qat_tosa_INT() -> None:
+    pipeline = TosaPipelineINT[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_op=[],
+        exir_op=[],
+        use_to_edge_transform_and_lower=True,
+        is_qat=True,
+        frobenius_threshold=None,
+        cosine_threshold=None,
+        qtol=12,
+    )
+    _set_nss_calibration_samples(pipeline)
+    pipeline.run()
+
+
+@common.SkipIfNoModelConverter
+def test_nss_qat_vgf_INT() -> None:
+    pipeline = VgfPipeline[input_t](
+        nss().eval(),
+        example_inputs(),
+        aten_op=[],
+        exir_op=[],
+        symmetric_io_quantization=True,
+        use_to_edge_transform_and_lower=True,
+        run_on_vulkan_runtime=True,
+        quantize=True,
+        is_qat=True,
+        tosa_version="TOSA-1.0+INT",
+        qtol=12,
+    )
+    _set_nss_calibration_samples(pipeline)
+    pipeline.run()

backends/arm/test/tester/test_pipeline.py

@@ -435,6 +435,7 @@ class TosaPipelineINT(TOSAPipeline, Generic[T]):
        run_on_tosa_ref_model: Set to true to test the tosa file on the TOSA reference model.
        symmetric_io_quantization: Whether to use symmetric I/O quantization.
        per_channel_quantization: Whether to use per-channel quantization.
+       is_qat: Whether to prepare the model for quantization-aware training instead of PTQ.
        use_to_edge_transform_and_lower: Selects betweeen two possible ways of lowering the module.
        custom_path : Path to dump intermediate artifacts such as tosa and pte to.
        tosa_debug_mode: Optional debug mode for TOSA compilation.
@@ -475,6 +476,7 @@ def __init__(
         tosa_extensions: Optional[List[str]] = None,
         epsilon: float = 2**-16,
         fold_quantize: bool = True,
+        is_qat: bool = False,
     ):
         if tosa_extensions is None:
             tosa_extensions = []
@@ -492,16 +494,22 @@ def __init__(
         # choose 16A8W quantization config when int16 extension is requested
         if "int16" in tosa_extensions:
             quantization_config = get_symmetric_a16w8_quantization_config(
-                is_per_channel=per_channel_quantization, epsilon=epsilon
+                is_per_channel=per_channel_quantization,
+                is_qat=is_qat,
+                epsilon=epsilon,
             )
         else:
             quantization_config = get_symmetric_quantization_config(
-                is_per_channel=per_channel_quantization
+                is_per_channel=per_channel_quantization,
+                is_qat=is_qat,
             )
         if symmetric_io_quantization:
             quantizer.set_io(quantization_config)
         quant_stage = Quantize(
-            quantizer, quantization_config, fold_quantize=fold_quantize
+            quantizer,
+            quantization_config,
+            is_qat=is_qat,
+            fold_quantize=fold_quantize,
         )
 
         super().__init__(
@@ -1217,6 +1225,7 @@ class VgfPipeline(BasePipeline, Generic[T]):
        run_on_vulkan_runtime: Whether to test VGF output on VKML runtime.
 
        vgf_compiler_flags: Optional compiler flags.
+       is_qat: Whether to prepare the model for quantization-aware training instead of PTQ.
 
        tosa_version: A string for identifying the TOSA version.
        tosa_spec: Optional override for the TOSA specification.
@@ -1255,7 +1264,11 @@ def __init__(
         fold_quantize: bool = True,
         preserve_io_quantization: bool = False,
         n_expected_delegates: int = 1,
+        is_qat: bool = False,
     ):
+        if is_qat and not quantize:
+            raise ValueError("QAT requires quantize=True.")
+
         if tosa_spec is None:
             if tosa_version is None:
                 tosa_version = str(VgfCompileSpec().tosa_spec)
@@ -1301,12 +1314,16 @@ def __init__(
         if quantize:
             quantizer = VgfQuantizer(compile_spec)
             quantization_config = get_symmetric_quantization_config(
-                is_per_channel=per_channel_quantization
+                is_per_channel=per_channel_quantization,
+                is_qat=is_qat,
             )
             if symmetric_io_quantization:
                 quantizer.set_io(quantization_config)
             quant_stage = Quantize(
-                quantizer, quantization_config, fold_quantize=fold_quantize
+                quantizer,
+                quantization_config,
+                is_qat=is_qat,
+                fold_quantize=fold_quantize,
             )
 
             self.add_stage(self.tester.quantize, quant_stage, pos=0)