[Misc] Cleanup: use consistent pytest-mock in unit tests (#2698)

Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>

tests/diffusion/offloader/test_sequential_backend.py

@@ -3,8 +3,6 @@
 
 """Unit tests for SequentialOffloadBackend."""
 
-from unittest.mock import patch
-
 import pytest
 import torch
 from torch import nn
@@ -44,7 +42,7 @@ def _track_pin_memory_calls():
 
 
 class TestMoveParamsPinMemory:
-    def test_dtensor_skips_pin_memory(self, accelerator_device):
+    def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """DTensor should skip pin_memory to avoid RuntimeError."""
         module = _create_simple_module().to(accelerator_device)
         tracker, mock_pin = _track_pin_memory_calls()
@@ -56,73 +54,73 @@ class TestMoveParamsPinMemory:
                 return True
             return original_isinstance(obj, cls)
 
-        with patch.object(torch.Tensor, "pin_memory", mock_pin):
-            with patch("builtins.isinstance", fake_isinstance):
-                hook = SequentialOffloadHook(
-                    offload_targets=[],
-                    device=accelerator_device,
-                    pin_memory=True,
-                    use_hsdp=False,
-                )
-                hook._move_params(
-                    module,
-                    torch.device("cpu"),
-                    non_blocking=False,
-                    pin_memory=True,
-                )
-                assert not tracker["called"], "pin_memory should not be called for DTensor"
+        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
+        monkeypatch.setattr("builtins.isinstance", fake_isinstance)
+        hook = SequentialOffloadHook(
+            offload_targets=[],
+            device=accelerator_device,
+            pin_memory=True,
+            use_hsdp=False,
+        )
+        hook._move_params(
+            module,
+            torch.device("cpu"),
+            non_blocking=False,
+            pin_memory=True,
+        )
+        assert not tracker["called"], "pin_memory should not be called for DTensor"
 
-    def test_regular_tensor_calls_pin_memory(self, accelerator_device):
+    def test_regular_tensor_calls_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """Regular tensor should call pin_memory when moving to CPU."""
         module = _create_simple_module().to(accelerator_device)
         tracker, mock_pin = _track_pin_memory_calls()
 
-        with patch.object(torch.Tensor, "pin_memory", mock_pin):
-            hook = SequentialOffloadHook(
-                offload_targets=[],
-                device=accelerator_device,
-                pin_memory=True,
-                use_hsdp=False,
-            )
-            hook._move_params(
-                module,
-                torch.device("cpu"),
-                non_blocking=False,
-                pin_memory=True,
-            )
-            assert tracker["called"], "pin_memory should be called for regular tensors"
+        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
+        hook = SequentialOffloadHook(
+            offload_targets=[],
+            device=accelerator_device,
+            pin_memory=True,
+            use_hsdp=False,
+        )
+        hook._move_params(
+            module,
+            torch.device("cpu"),
+            non_blocking=False,
+            pin_memory=True,
+        )
+        assert tracker["called"], "pin_memory should be called for regular tensors"
 
-    def test_pin_memory_skipped_when_disabled(self, accelerator_device):
+    def test_pin_memory_skipped_when_disabled(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """pin_memory should not be called when pin_memory=False."""
         module = _create_simple_module().to(accelerator_device)
         tracker, mock_pin = _track_pin_memory_calls()
 
-        with patch.object(torch.Tensor, "pin_memory", mock_pin):
-            hook = SequentialOffloadHook(
-                offload_targets=[],
-                device=accelerator_device,
-                pin_memory=False,
-                use_hsdp=False,
-            )
-            hook._move_params(
-                module,
-                torch.device("cpu"),
-                non_blocking=False,
-                pin_memory=False,
-            )
-            assert not tracker["called"], "pin_memory should not be called when disabled"
+        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
+        hook = SequentialOffloadHook(
+            offload_targets=[],
+            device=accelerator_device,
+            pin_memory=False,
+            use_hsdp=False,
+        )
+        hook._move_params(
+            module,
+            torch.device("cpu"),
+            non_blocking=False,
+            pin_memory=False,
+        )
+        assert not tracker["called"], "pin_memory should not be called when disabled"
 
-    def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device):
+    def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
         """pin_memory should not be called for non-CPU targets."""
         module = _create_simple_module().to("cpu")
         tracker, mock_pin = _track_pin_memory_calls()
 
-        with patch.object(torch.Tensor, "pin_memory", mock_pin):
-            hook = SequentialOffloadHook(
-                offload_targets=[],
-                device=torch.device("cpu"),
-                pin_memory=True,
-                use_hsdp=False,
-            )
-            hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True)
-            assert not tracker["called"], "pin_memory should not be called for non-CPU target"
+        monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
+        hook = SequentialOffloadHook(
+            offload_targets=[],
+            device=torch.device("cpu"),
+            pin_memory=True,
+            use_hsdp=False,
+        )
+        hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True)
+        assert not tracker["called"], "pin_memory should not be called for non-CPU target"