add test cases

Signed-off-by: Hollow Man <[email protected]>

Author: HollowMan6

tests/unit_tests/peft/test_adapter_wrapper.py

@@ -25,7 +25,9 @@
 import torch
 import torch.nn as nn
 
+from megatron.bridge.peft.base import PEFT
 from megatron.bridge.peft.adapter_wrapper import AdapterWrapper
+from megatron.bridge.peft.lora_layers import LoRALinear
 
 
 class MockLinear(nn.Module):
@@ -67,6 +69,25 @@ def forward(self, x, *args, **kwargs):
         return linear_output + adapter_output, bias
 
 
+class DummyPEFT(PEFT):
+    """Minimal PEFT implementation for adapter enable/disable tests."""
+
+    def transform(self, module: nn.Module, name=None, prefix=None) -> nn.Module:
+        return module
+
+
+class AdapterModel(nn.Module):
+    """Model with a single LoRALinear adapter wrapper for testing."""
+
+    def __init__(self, to_wrap: nn.Module, adapter: nn.Module) -> None:
+        super().__init__()
+        self.lora = LoRALinear(to_wrap, adapter)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output, _ = self.lora(x)
+        return output
+
+
 class TestAdapterWrapper:
     """Test the AdapterWrapper base class."""
 
@@ -258,3 +279,59 @@ def test_adapter_wrapper_is_abstract(self):
         assert hasattr(AdapterWrapper, "base_linear_forward")
         assert hasattr(AdapterWrapper, "state_dict")
         assert hasattr(AdapterWrapper, "sharded_state_dict")
+
+    def test_adapter_wrapper_enable_disable_toggle(self, mock_linear_simple, simple_adapter):
+        """Test adapter output toggling via AdapterWrapper methods."""
+        wrapper = LoRALinear(mock_linear_simple, simple_adapter)
+        x = torch.randn(5, 10)
+
+        base_output, _ = mock_linear_simple(x)
+        enabled_output, _ = wrapper(x)
+        expected = base_output + simple_adapter(x)
+        assert torch.allclose(enabled_output, expected, atol=1e-6)
+
+        wrapper.disable_adapter_layers()
+        disabled_output, _ = wrapper(x)
+        assert torch.allclose(disabled_output, base_output, atol=1e-6)
+
+        wrapper.enable_adapter_layers()
+        reenabled_output, _ = wrapper(x)
+        assert torch.allclose(reenabled_output, enabled_output, atol=1e-6)
+
+    def test_peft_disable_adapter_context_manager(self, mock_linear_simple, simple_adapter):
+        """Test PEFT.disable_adapter restores adapter state."""
+        peft = DummyPEFT()
+        model = AdapterModel(mock_linear_simple, simple_adapter)
+        x = torch.randn(5, 10)
+
+        base_output, _ = mock_linear_simple(x)
+        enabled_output = model(x)
+
+        with peft.disable_adapter(model):
+            disabled_output = model(x)
+            assert torch.allclose(disabled_output, base_output, atol=1e-6)
+
+        assert torch.allclose(model(x), enabled_output, atol=1e-6)
+
+        with pytest.raises(RuntimeError):
+            with peft.disable_adapter(model):
+                raise RuntimeError("boom")
+
+        assert torch.allclose(model(x), enabled_output, atol=1e-6)
+
+    def test_peft_enable_disable_adapter_layers_manual(self, mock_linear_simple, simple_adapter):
+        """Test manual adapter enable/disable via PEFT helpers."""
+        peft = DummyPEFT()
+        model = AdapterModel(mock_linear_simple, simple_adapter)
+        x = torch.randn(5, 10)
+
+        base_output, _ = mock_linear_simple(x)
+        enabled_output = model(x)
+
+        peft.disable_adapter_layers(model)
+        disabled_output = model(x)
+        assert torch.allclose(disabled_output, base_output, atol=1e-6)
+
+        peft.enable_adapter_layers(model)
+        reenabled_output = model(x)
+        assert torch.allclose(reenabled_output, enabled_output, atol=1e-6)