[lora] feat: support disabling dense adapters for reference inference

docs/key_features/lora.md

@@ -727,3 +727,21 @@ torchrun --nproc_per_node=4 tests/lora/test_moe_lora_trainer.py \
 1. Writer trains `max_steps=4`, writes DCP shards at step 2 / step 4 and the HF LoRA adapter (`adapter_config.json` with the `veomni_lora` block + `adapter_model.safetensors`, no sidecar) at the same cadence; snapshots the full LoRA tensors at `on_train_begin` and `on_train_end`.
 2. DCP-resume subprocess loads the step-2 DCP shard, continues to step 4, and the resulting LoRA tensors must be **bit-exact** vs the writer's `on_train_end` snapshot (DCP ships model + optimizer + RNG + dataloader state).
 3. LoRA-adapter-resume subprocess loads the step-4 HF adapter via `model.lora_config.lora_adapter`, and the resulting `on_train_begin` snapshot (= post-load, pre-train) must be bit-exact vs the writer's `on_train_end` snapshot in bf16 (the on-disk adapter dtype).
+
+### Temporarily disabling native dense adapters
+
+`VeOmniLoraModel.disable_adapter()` bypasses dense adapter contributions without
+changing base weights, adapter parameters, or their `requires_grad` flags:
+
+```python
+with model.disable_adapter():
+    with torch.no_grad():
+        reference_output = model(input_ids=input_ids)
+```
+
+The context restores the previous state on exit, including nested contexts and
+exceptions. It requires unmerged dense adapters; expert adapters are not yet
+supported. To use the result as a frozen base reference policy, train with
+`bias="none"` and no extra trainable base modules. Disabling adapters does not
+undo updates to trainable base biases or other base parameters. Keep the context
+active for any backward recomputation if using it outside inference.

scripts/ci/accelerator/suites/gpu.json

@@ -604,6 +604,16 @@
       "tests/utils/test_profiler_memory.py"
     ]
   },
+  {
+    "name": "native-dense-lora-unit",
+    "suite": "unit",
+    "conditional": false,
+    "args": [
+      "-s",
+      "-x",
+      "tests/lora/test_veomni_lora_native.py"
+    ]
+  },
   {
     "name": "58-unit",
     "suite": "unit",

scripts/ci/accelerator/suites/npu.json

@@ -407,6 +407,16 @@
       "tests/utils/test_moe_router_replay_api.py"
     ]
   },
+  {
+    "name": "native-dense-lora-unit",
+    "suite": "unit",
+    "conditional": false,
+    "args": [
+      "-s",
+      "-x",
+      "tests/lora/test_veomni_lora_native.py"
+    ]
+  },
   {
     "name": "npu-unit-39",
     "suite": "unit",

tests/lora/test_veomni_lora_native.py

@@ -11,9 +11,9 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""CPU single-process tests for the VeOmni-native dense LoRA stack (Phase 1).
+"""Single-process tests for the VeOmni-native dense LoRA stack (Phase 1).
 
-No distributed / GPU / HF-model build required — a tiny ``nn.Module`` exercises:
+The CPU cases require no distributed / GPU / HF-model build; a tiny ``nn.Module`` exercises:
 
 * ``VeOmniLoraConfig`` YAML <-> config <-> ``adapter_config.json`` round-trips.
 * Dense injection: PEFT-identical FQNs, trainable-only-LoRA, no-op-at-init.
@@ -24,6 +24,7 @@
 
 ``test_peft_bidirectional_interop`` additionally asserts cross-compatibility
 with stock ``peft`` when it is installed (dev/test-only dependency).
+The adapter-bypass regression also runs on the runner-selected accelerator.
 """
 
 from __future__ import annotations
@@ -35,11 +36,16 @@
 import torch
 import torch.nn as nn
 
+from scripts.ci.accelerator.launch_plan import LaunchSpec
 from tests.tools.training_utils import unbuilt_runtime
 from veomni.lora import VeOmniLoraConfig, VeOmniLoraModel
 from veomni.lora.layers import LoraLinear
 from veomni.lora.state_dict import get_lora_state_dict, load_adapter_state_dict
 from veomni.lora.weight_loading import load_lora_weights
+from veomni.utils.device import get_device_type
+
+
+pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu"))
 
 
 torch.manual_seed(0)
@@ -303,3 +309,45 @@ def test_peft_bidirectional_interop(tmp_path):
 
 if __name__ == "__main__":
     raise SystemExit(pytest.main([__file__, "-v"]))
+
+
+@pytest.mark.parametrize(
+    "device_kind",
+    [
+        pytest.param("cpu", id="cpu"),
+        pytest.param(
+            "accelerator", id="accelerator", marks=pytest.mark.launch_plan(provider=LaunchSpec(device="accelerator"))
+        ),
+    ],
+)
+def test_disable_adapter_restores_nested_context_after_exception(device_kind):
+    """Reference inference bypasses nonzero adapters without changing weights."""
+    device = "cpu" if device_kind == "cpu" else get_device_type()
+    if device_kind == "accelerator" and device == "cpu":
+        pytest.skip("Adapter bypass accelerator coverage requires an available accelerator")
+    base = Toy().to(device)
+    reference = copy.deepcopy(base)
+    model = VeOmniLoraModel(base, _base_config()).to(device)
+    _randomize_lora_b(model)
+    x = torch.randn(3, 32, device=device)
+    adapted = model(x).detach().clone()
+    expected = reference(x).detach()
+    assert not torch.equal(adapted, expected)
+    with pytest.raises(RuntimeError, match="test exception"):
+        with model.disable_adapter():
+            torch.testing.assert_close(model(x), expected, rtol=0, atol=0)
+            with model.disable_adapter():
+                torch.testing.assert_close(model(x), expected, rtol=0, atol=0)
+            torch.testing.assert_close(model(x), expected, rtol=0, atol=0)
+            raise RuntimeError("test exception")
+    torch.testing.assert_close(model(x), adapted, rtol=0, atol=0)
+
+
+def test_disable_adapter_rejects_merged_layers_before_changing_flags():
+    model = VeOmniLoraModel(Toy(), _base_config())
+    layers = [m for m in model.modules() if isinstance(m, LoraLinear)]
+    layers[-1].merge()
+    with pytest.raises(ValueError, match="unmerged"):
+        with model.disable_adapter():
+            pytest.fail("Merged weights cannot provide the base-policy reference")
+    assert all(not layer.disable_adapters for layer in layers)

veomni/lora/layers.py

@@ -51,6 +51,7 @@ def _init_lora_layer(self, base_layer: nn.Module) -> None:
         self.lora_alpha: dict[str, int] = {}
         self.use_rslora: dict[str, bool] = {}
         self._merged_adapters: list[str] = []
+        self.disable_adapters = False
 
     @property
     def merged(self) -> bool:
@@ -141,7 +142,7 @@ def __init__(
 
     def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         result = self.base_layer(x, *args, **kwargs)
-        if self.merged:
+        if self.merged or self.disable_adapters:
             return result
         adapter = self.active_adapter
         if adapter not in self.lora_A:

veomni/lora/model.py

@@ -27,6 +27,7 @@
 
 from __future__ import annotations
 
+from contextlib import contextmanager
 from typing import Any
 
 import torch.nn as nn
@@ -131,6 +132,30 @@ def get_lora_config(self, adapter_name: str | None = None) -> VeOmniLoraConfig:
     def get_base_model(self) -> nn.Module:
         return self.base_model.model
 
+    @contextmanager
+    def disable_adapter(self):
+        """Temporarily bypass dense adapters for reference-policy inference.
+
+        Restores every layer's previous state, including on nested use or an
+        exception. Merged weights and expert adapters require separate handling
+        and are rejected before any layer is changed.
+        """
+        from .layers import LoraLinear
+
+        if self.get_lora_config().has_moe:
+            raise NotImplementedError("disable_adapter() currently supports dense LoRA only.")
+        layers = [module for module in self.modules() if isinstance(module, LoraLinear)]
+        if any(layer.merged for layer in layers):
+            raise ValueError("disable_adapter() requires unmerged LoRA layers.")
+        previous = [layer.disable_adapters for layer in layers]
+        try:
+            for layer in layers:
+                layer.disable_adapters = True
+            yield
+        finally:
+            for layer, disabled in zip(layers, previous):
+                layer.disable_adapters = disabled
+
     def forward(self, *args: Any, **kwargs: Any) -> Any:
         return self.base_model(*args, **kwargs)