From 7a5779ff83a847b75bfe552de1cf6d62afd34a9b Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Wed, 21 Oct 2026 15:28:09 +0800 Subject: [PATCH 1/2] [lora] feat: allow temporarily disabling dense adapters --- .github/workflows/gpu_unit_tests.yml | 1 + .github/workflows/npu_unit_tests.yml | 1 + docs/key_features/lora.md | 18 ++++++++++++++++ tests/lora/test_veomni_lora_native.py | 30 +++++++++++++++++++++++++++ veomni/lora/layers.py | 3 ++- veomni/lora/model.py | 25 ++++++++++++++++++++++ 6 files changed, 77 insertions(+), 1 deletion(-) diff --git a/.github/workflows/gpu_unit_tests.yml b/.github/workflows/gpu_unit_tests.yml index f33c3ff83d..58578ab421 100644 --- a/.github/workflows/gpu_unit_tests.yml +++ b/.github/workflows/gpu_unit_tests.yml @@ -231,6 +231,7 @@ jobs: # Order: cheap unit tests first to fail fast, trainer/ep2 last (multi-subprocess, ~5-10 min each). run: | uv run --frozen pytest -s -x tests/lora/test_target_mapping.py + uv run --frozen pytest -s -x tests/lora/test_veomni_lora_native.py uv run --frozen pytest -s -x tests/lora/test_qwen3_5_moe_lora.py uv run --frozen pytest -s -x tests/lora/test_dense_mlp_lora.py uv run --frozen pytest -s -x tests/lora/test_veomni_lora_moe_native.py diff --git a/.github/workflows/npu_unit_tests.yml b/.github/workflows/npu_unit_tests.yml index 3b676b9654..cec202abb7 100644 --- a/.github/workflows/npu_unit_tests.yml +++ b/.github/workflows/npu_unit_tests.yml @@ -247,6 +247,7 @@ jobs: # sync needed here. run: | uv run --frozen pytest -v -s -x tests/lora/test_lora_trainer_saveload.py + uv run --frozen pytest -v -s -x tests/lora/test_veomni_lora_native.py uv run --frozen pytest -v -s -x tests/lora/test_qwen3_5_moe_lora.py - name: Run MoE-LoRA ep_sharded_stream_load (PEFT, EP=2) # Exercises the PEFT ep_sharded adapter streaming *and* the NPU diff --git a/docs/key_features/lora.md b/docs/key_features/lora.md index 7066fb258b..c07475c5be 100644 --- a/docs/key_features/lora.md +++ b/docs/key_features/lora.md @@ -723,3 +723,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/tests/lora/test_veomni_lora_native.py b/tests/lora/test_veomni_lora_native.py index 96b7838b92..030ba03f2e 100644 --- a/tests/lora/test_veomni_lora_native.py +++ b/tests/lora/test_veomni_lora_native.py @@ -303,3 +303,33 @@ def test_peft_bidirectional_interop(tmp_path): if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"])) + + +def test_disable_adapter_restores_nested_context_after_exception(): + """Reference inference bypasses nonzero adapters without changing weights.""" + base = Toy() + reference = copy.deepcopy(base) + model = VeOmniLoraModel(base, _base_config()) + _randomize_lora_b(model) + x = torch.randn(3, 32) + 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) From 0440ba983b063d02fffc581faad09316fcb35ad7 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Fri, 23 Oct 2026 13:38:17 +0800 Subject: [PATCH 2/2] [lora, ci] test: cover adapter bypass on the selected accelerator --- tests/lora/test_veomni_lora_native.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/tests/lora/test_veomni_lora_native.py b/tests/lora/test_veomni_lora_native.py index 0041d263bc..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 @@ -41,6 +42,7 @@ 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")) @@ -309,13 +311,25 @@ def test_peft_bidirectional_interop(tmp_path): raise SystemExit(pytest.main([__file__, "-v"])) -def test_disable_adapter_restores_nested_context_after_exception(): +@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.""" - base = Toy() + 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()) + model = VeOmniLoraModel(base, _base_config()).to(device) _randomize_lora_b(model) - x = torch.randn(3, 32) + x = torch.randn(3, 32, device=device) adapted = model(x).detach().clone() expected = reference(x).detach() assert not torch.equal(adapted, expected)