diff --git a/docs/key_features/lora.md b/docs/key_features/lora.md index e23381e98f..a68d58cc9d 100644 --- a/docs/key_features/lora.md +++ b/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. diff --git a/scripts/ci/accelerator/suites/gpu.json b/scripts/ci/accelerator/suites/gpu.json index 422591c7bb..7d7ba4a75d 100644 --- a/scripts/ci/accelerator/suites/gpu.json +++ b/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", diff --git a/scripts/ci/accelerator/suites/npu.json b/scripts/ci/accelerator/suites/npu.json index e8644c5baf..5be2c07e85 100644 --- a/scripts/ci/accelerator/suites/npu.json +++ b/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", diff --git a/tests/lora/test_veomni_lora_native.py b/tests/lora/test_veomni_lora_native.py index 96b7838b92..9cc7e3e6be 100644 --- a/tests/lora/test_veomni_lora_native.py +++ b/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) diff --git a/veomni/lora/layers.py b/veomni/lora/layers.py index 1087c178c8..c71ad44342 100644 --- a/veomni/lora/layers.py +++ b/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: diff --git a/veomni/lora/model.py b/veomni/lora/model.py index b3f146dcd5..e1ad4de218 100644 --- a/veomni/lora/model.py +++ b/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)