memory

Author: sayakpaul

tests/models/testing_utils/common.py

@@ -262,7 +262,39 @@ def get_dummy_inputs(self) -> Dict[str, Any]:
         raise NotImplementedError("Subclasses must implement `get_dummy_inputs()`.")
 
 
-class ModelTesterMixin:
+class BaseModelOutputMixin:
+    """Provides the class-scoped `base_model_output` fixture shared across tester mixins.
+
+    Kept separate from `BaseModelTesterConfig` — which only declares the testing contract and performs no
+    computation — so any mixin that needs the cached reference output (`ModelTesterMixin`, the memory
+    offload mixins, ...) can inherit it without duplicating the build-and-forward.
+    """
+
+    @pytest.fixture(scope="class")
+    def base_model_output(self):
+        """Class-scoped reference forward output, built once and reused across the class.
+
+        Building the model and running its forward pass is fully deterministic (`torch.manual_seed(0)`
+        plus the deterministic `get_dummy_inputs` contract), so the reference ("base") output is
+        identical for every test in the class. The save/load, parallelism, and memory-offload tests
+        compare a reloaded/offloaded model against this output; computing it a single time here — instead
+        of rebuilding the model and re-running the forward in each test — removes that redundant work and
+        speeds up the suite.
+
+        The hardware-gated tests that consume this fixture use `pytest.mark.skipif` (via the `require_*`
+        decorators), which pytest evaluates before fixture setup, so skipping on a machine without the
+        required accelerators never triggers this forward.
+
+        Tests that still need a live model (e.g. to save or offload it) build their own with the same
+        seed, so the reloaded model's weights match this cached output.
+        """
+        torch.manual_seed(0)
+        model = self.model_class(**self.get_init_dict()).eval().to(torch_device)
+        with torch.no_grad():
+            return model(**self.get_dummy_inputs(), return_dict=False)[0]
+
+
+class ModelTesterMixin(BaseModelOutputMixin):
     """
     Base mixin class for model testing with common test methods.
 
@@ -282,28 +314,6 @@ class TestMyModel(MyModelTestConfig, ModelTesterMixin):
             pass
     """
 
-    @pytest.fixture(scope="class")
-    def base_model_output(self):
-        """Class-scoped reference forward output, built once and reused across the class.
-
-        Building the model and running its forward pass is fully deterministic (`torch.manual_seed(0)`
-        plus the deterministic `get_dummy_inputs` contract), so the reference ("base") output is
-        identical for every test in the class. The save/load and parallelism tests compare a reloaded
-        model against this output; computing it a single time here — instead of rebuilding the model and
-        re-running the forward in each test — removes that redundant work and speeds up the suite.
-
-        The hardware-gated tests that consume this fixture use `pytest.mark.skipif` (via the
-        `require_*` decorators), which pytest evaluates before fixture setup, so skipping on a machine
-        without the required accelerators never triggers this forward.
-
-        Tests that still need a live model (e.g. to save it) build their own with the same seed, so the
-        reloaded model's weights match this cached output.
-        """
-        torch.manual_seed(0)
-        model = self.model_class(**self.get_init_dict()).eval().to(torch_device)
-        with torch.no_grad():
-            return model(**self.get_dummy_inputs(), return_dict=False)[0]
-
     @torch.no_grad()
     def test_from_save_pretrained(self, base_model_output, tmp_path, atol=5e-5, rtol=5e-5):
         torch.manual_seed(0)

tests/models/testing_utils/memory.py

@@ -37,7 +37,7 @@
     require_accelerator,
     torch_device,
 )
-from .common import cast_inputs_to_dtype, check_device_map_is_respected
+from .common import BaseModelOutputMixin, cast_inputs_to_dtype, check_device_map_is_respected
 
 
 def require_offload_support(func):
@@ -69,7 +69,7 @@ def wrapper(self, *args, **kwargs):
 
 
 @is_cpu_offload
-class CPUOffloadTesterMixin:
+class CPUOffloadTesterMixin(BaseModelOutputMixin):
     """
     Mixin class for testing CPU offloading functionality.
 
@@ -94,16 +94,14 @@ def model_split_percents(self) -> list[float]:
 
     @require_offload_support
     @torch.no_grad()
-    def test_cpu_offload(self, tmp_path, atol=1e-5, rtol=0):
+    def test_cpu_offload(self, base_model_output, tmp_path, atol=1e-5, rtol=0):
+        torch.manual_seed(0)
         config = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**config).eval()
 
         model = model.to(torch_device)
 
-        torch.manual_seed(0)
-        base_output = model(**inputs_dict)
-
         model_size = compute_module_sizes(model)[""]
         # We test several splits of sizes to make sure it works
         max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
@@ -120,21 +118,19 @@ def test_cpu_offload(self, tmp_path, atol=1e-5, rtol=0):
             new_output = new_model(**inputs_dict)
 
             assert_tensors_close(
-                base_output[0], new_output[0], atol=atol, rtol=rtol, msg="Output should match with CPU offloading"
+                base_model_output, new_output[0], atol=atol, rtol=rtol, msg="Output should match with CPU offloading"
             )
 
     @require_offload_support
     @torch.no_grad()
-    def test_disk_offload_without_safetensors(self, tmp_path, atol=1e-5, rtol=0):
+    def test_disk_offload_without_safetensors(self, base_model_output, tmp_path, atol=1e-5, rtol=0):
+        torch.manual_seed(0)
         config = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**config).eval()
 
         model = model.to(torch_device)
 
-        torch.manual_seed(0)
-        base_output = model(**inputs_dict)
-
         model_size = compute_module_sizes(model)[""]
         max_size = int(self.model_split_percents[0] * model_size)
         # Force disk offload by setting very small CPU memory
@@ -154,21 +150,19 @@ def test_disk_offload_without_safetensors(self, tmp_path, atol=1e-5, rtol=0):
         new_output = new_model(**inputs_dict)
 
         assert_tensors_close(
-            base_output[0], new_output[0], atol=atol, rtol=rtol, msg="Output should match with disk offloading"
+            base_model_output, new_output[0], atol=atol, rtol=rtol, msg="Output should match with disk offloading"
         )
 
     @require_offload_support
     @torch.no_grad()
-    def test_disk_offload_with_safetensors(self, tmp_path, atol=1e-5, rtol=0):
+    def test_disk_offload_with_safetensors(self, base_model_output, tmp_path, atol=1e-5, rtol=0):
+        torch.manual_seed(0)
         config = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
         model = self.model_class(**config).eval()
 
         model = model.to(torch_device)
 
-        torch.manual_seed(0)
-        base_output = model(**inputs_dict)
-
         model_size = compute_module_sizes(model)[""]
         model.cpu().save_pretrained(str(tmp_path))
 
@@ -183,7 +177,7 @@ def test_disk_offload_with_safetensors(self, tmp_path, atol=1e-5, rtol=0):
         new_output = new_model(**inputs_dict)
 
         assert_tensors_close(
-            base_output[0],
+            base_model_output,
             new_output[0],
             atol=atol,
             rtol=rtol,
@@ -192,7 +186,7 @@ def test_disk_offload_with_safetensors(self, tmp_path, atol=1e-5, rtol=0):
 
 
 @is_group_offload
-class GroupOffloadTesterMixin:
+class GroupOffloadTesterMixin(BaseModelOutputMixin):
     """
     Mixin class for testing group offloading functionality.
 
@@ -209,10 +203,9 @@ class GroupOffloadTesterMixin:
 
     @require_group_offload_support
     @pytest.mark.parametrize("record_stream", [False, True])
-    def test_group_offloading(self, record_stream, atol=1e-5, rtol=0):
+    def test_group_offloading(self, base_model_output, record_stream, atol=1e-5, rtol=0):
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
-        torch.manual_seed(0)
 
         @torch.no_grad()
         def run_forward(model):
@@ -224,10 +217,7 @@ def run_forward(model):
             model.eval()
             return model(**inputs_dict)[0]
 
-        model = self.model_class(**init_dict)
-
-        model.to(torch_device)
-        output_without_group_offloading = run_forward(model)
+        output_without_group_offloading = base_model_output
 
         torch.manual_seed(0)
         model = self.model_class(**init_dict)