From aca9ec27c7238005b3e6210cf1afd9ee7cfac424 Mon Sep 17 00:00:00 2001 From: Jenny Luo Date: Wed, 14 Oct 2026 12:00:48 -0700 Subject: [PATCH 1/3] lora: honor lora_target_modules and fix the sidecar save crash training.lora_target_modules was written to the adapter sidecar but never passed to apply_lora_torchtune, so every nn.Linear except lm_head was adapted regardless of the setting. Setting it via Hydra also crashed save_lora_adapter, because json cannot serialize an OmegaConf ListConfig. apply_lora_torchtune now takes target_modules (leaf names, matched like exclude) and raises on a name that matches no layer, an empty list, or a bare string. ZoneATrainer and ZoneDTrainer pass the config value through. With target_modules unset the result is unchanged (state dict bit-identical on a seeded toy backbone). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GmpnEoNqJpQjR59ie7ZMek --- src/conf/training/bioreason_sft.yaml | 2 +- src/config.py | 4 +- src/training/trainer_grpo.py | 1 + src/training/trainer_zone_a.py | 4 +- src/utils/lora_utils.py | 40 +++++++++++++- tests/test_lora_target_modules.py | 82 ++++++++++++++++++++++++++++ 6 files changed, 127 insertions(+), 6 deletions(-) create mode 100644 tests/test_lora_target_modules.py diff --git a/src/conf/training/bioreason_sft.yaml b/src/conf/training/bioreason_sft.yaml index a439b8d..4a33c47 100644 --- a/src/conf/training/bioreason_sft.yaml +++ b/src/conf/training/bioreason_sft.yaml @@ -72,7 +72,7 @@ lora_enabled: true lora_r: 32 lora_alpha: 64 lora_dropout: 0.05 -# lora_target_modules: null # null = auto-detect all nn.Linear layers (excluding lm_head) +# lora_target_modules: null # null = all nn.Linear layers (excluding lm_head); or leaf names, e.g. [q_proj, v_proj] # Data format: include reasoning traces in Stage 2 so LoRA learns chain-of-thought. use_reasoning_traces: true diff --git a/src/config.py b/src/config.py index bbc8a6a..bad57eb 100644 --- a/src/config.py +++ b/src/config.py @@ -146,7 +146,9 @@ class TrainingConfig: lora_r: int = 32 lora_alpha: int = 64 lora_dropout: float = 0.05 - lora_target_modules: list[str] | None = None # Auto-detected if None + # Leaf module names to adapt, e.g. ["q_proj", "v_proj"]; None adapts every + # nn.Linear except lm_head. + lora_target_modules: list[str] | None = None # Alpha to use when merging a resumed SFT-stage LoRA adapter into the base # weights before creating fresh GRPO-rank LoRALinear layers (the two # stages' adapters aren't rank-compatible for direct loading). Defaults to diff --git a/src/training/trainer_grpo.py b/src/training/trainer_grpo.py index 68e884b..d11bdfe 100644 --- a/src/training/trainer_grpo.py +++ b/src/training/trainer_grpo.py @@ -231,6 +231,7 @@ def __init__( rank=config.lora_r, alpha=config.lora_alpha, dropout=config.lora_dropout, + target_modules=getattr(config, "lora_target_modules", None), ) # Verify the freshly created LoRA layers are actually rank config.lora_r diff --git a/src/training/trainer_zone_a.py b/src/training/trainer_zone_a.py index 22ed975..c626ece 100644 --- a/src/training/trainer_zone_a.py +++ b/src/training/trainer_zone_a.py @@ -273,7 +273,8 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load param.requires_grad = True if lora_enabled: - # torchtune replaces every nn.Linear (except lm_head) with LoRALinear + # torchtune replaces every nn.Linear (except lm_head), or only the + # leaf names in lora_target_modules when set, with LoRALinear # in-place, copies pretrained weights, then freezes base weights via # set_trainable_params — only lora_a.weight / lora_b.weight keep # requires_grad=True afterward. Same mechanism trainer_grpo.py @@ -285,6 +286,7 @@ def __init__(self, model: UnifiedTransformer, config: TrainingConfig, train_load rank=config.lora_r, alpha=config.lora_alpha, dropout=config.lora_dropout, + target_modules=getattr(config, "lora_target_modules", None), ) self.accelerator.print("torchtune LoRA applied to backbone.") elif not freeze_llm: diff --git a/src/utils/lora_utils.py b/src/utils/lora_utils.py index 2a890b4..540cff5 100644 --- a/src/utils/lora_utils.py +++ b/src/utils/lora_utils.py @@ -31,8 +31,12 @@ def apply_lora_torchtune( alpha: float = 64.0, dropout: float = 0.05, exclude: tuple[str, ...] = ("lm_head",), + target_modules: list[str] | None = None, ) -> nn.Module: - """Replace every nn.Linear in backbone (except exclusions) with torchtune LoRALinear. + """Replace nn.Linear layers in backbone with torchtune LoRALinear. + + By default every nn.Linear except the exclusions is replaced; target_modules + restricts replacement to the listed leaf names. Base weights are copied from the original layer, then frozen via set_trainable_params. Only lora_a.weight and lora_b.weight remain trainable. @@ -43,12 +47,28 @@ def apply_lora_torchtune( alpha: LoRA alpha scaling factor. dropout: Dropout applied to the LoRA path. exclude: Leaf module names to skip (default: lm_head). + target_modules: Leaf module names to adapt, e.g. ["q_proj", "v_proj"]. + None (default) adapts every nn.Linear not in exclude. Returns: The modified backbone with LoRA layers. + + Raises: + TypeError: If target_modules is a bare string rather than a list of names. + ValueError: If target_modules is empty, or names a module that matches no + non-excluded nn.Linear (a typo would otherwise train nothing there). """ from torchtune.modules.peft import LoRALinear, get_adapter_params, set_trainable_params + if isinstance(target_modules, str): + raise TypeError( + f"target_modules must be a list of leaf module names, got the string " + f"{target_modules!r}; write [{target_modules!r}]" + ) + targets = None if target_modules is None else set(target_modules) + if targets is not None and not targets: + raise ValueError("target_modules is empty; use None to adapt every nn.Linear") + def _get_parent(model: nn.Module, dotted_name: str): parts = dotted_name.split(".") parent = model @@ -61,9 +81,17 @@ def _get_parent(model: nn.Module, dotted_name: str): for name, module in backbone.named_modules(): if isinstance(module, nn.Linear): leaf = name.split(".")[-1] - if leaf not in exclude: + if leaf not in exclude and (targets is None or leaf in targets): to_replace.append((name, module)) + if targets is not None: + unmatched = sorted(targets - {name.split(".")[-1] for name, _ in to_replace}) + if unmatched: + raise ValueError( + f"target_modules {unmatched} matched no nn.Linear in the backbone " + f"(excluded: {list(exclude)})" + ) + replaced = 0 for name, module in to_replace: parent, leaf = _get_parent(backbone, name) @@ -123,7 +151,13 @@ def save_lora_adapter( logger.info(f"LoRA adapter saved to {path} ({len(adapter_state)} tensors, {n:,} params)") if rank is not None or alpha is not None or target_modules is not None: - meta = {"rank": rank, "alpha": alpha, "target_modules": target_modules} + # list(): Hydra hands target_modules over as an OmegaConf ListConfig, + # which json cannot serialize. + meta = { + "rank": rank, + "alpha": alpha, + "target_modules": None if target_modules is None else list(target_modules), + } meta_path = _meta_path(path) with open(meta_path, "w") as f: json.dump(meta, f) diff --git a/tests/test_lora_target_modules.py b/tests/test_lora_target_modules.py new file mode 100644 index 0000000..f263dfb --- /dev/null +++ b/tests/test_lora_target_modules.py @@ -0,0 +1,82 @@ +"""apply_lora_torchtune honors target_modules, and the adapter sidecar records it. + +Previously ``training.lora_target_modules`` was written to the adapter sidecar +but never reached ``apply_lora_torchtune``, so every nn.Linear except lm_head was +adapted regardless; and a Hydra-supplied list crashed ``save_lora_adapter``'s +json.dump. Skipped where torchtune is not installed (CI does not install it). +""" + +from __future__ import annotations + +import json + +import pytest +import torch.nn as nn + +pytest.importorskip("torchtune") +from src.utils.lora_utils import apply_lora_torchtune, save_lora_adapter # noqa: E402 +from torchtune.modules.peft import LoRALinear # noqa: E402 + +pytestmark = pytest.mark.unit + + +class _Block(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(8, 8) + self.v_proj = nn.Linear(8, 8) + self.o_proj = nn.Linear(8, 8) + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Block(), _Block()]) + self.lm_head = nn.Linear(8, 16) + + +def _adapted(model: nn.Module) -> set[str]: + return {name for name, module in model.named_modules() if isinstance(module, LoRALinear)} + + +def test_default_adapts_every_linear_except_lm_head(): + model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0) + assert _adapted(model) == { + f"layers.{i}.{leaf}" for i in range(2) for leaf in ("q_proj", "v_proj", "o_proj") + } + + +def test_target_modules_restricts_adaptation_and_trainable_params(): + model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=["q_proj", "v_proj"]) + assert _adapted(model) == {f"layers.{i}.{leaf}" for i in range(2) for leaf in ("q_proj", "v_proj")} + trainable = {name for name, p in model.named_parameters() if p.requires_grad} + assert trainable == { + f"layers.{i}.{leaf}.lora_{ab}.weight" + for i in range(2) + for leaf in ("q_proj", "v_proj") + for ab in ("a", "b") + } + + +@pytest.mark.parametrize( + ("target_modules", "error"), + [ + (["q_prj"], ValueError), # typo matches nothing + (["lm_head"], ValueError), # excluded + ([], ValueError), + ("q_proj", TypeError), # bare string, not a list + ], +) +def test_bad_target_modules_raise(target_modules, error): + with pytest.raises(error): + apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=target_modules) + + +def test_sidecar_records_hydra_list(tmp_path): + omegaconf = pytest.importorskip("omegaconf") + targets = omegaconf.OmegaConf.create({"t": ["q_proj"]}).t # as cfg.training.get() returns it + model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=targets) + path = tmp_path / "lora_adapter.pt" + save_lora_adapter(model, str(path), rank=2, alpha=4.0, target_modules=targets) + meta = json.loads((tmp_path / "lora_adapter.pt.meta.json").read_text()) + assert meta == {"rank": 2, "alpha": 4.0, "target_modules": ["q_proj"]} From 5c9e8c2e8b1077f5e6611621a48d7a53c3d0a25b Mon Sep 17 00:00:00 2001 From: Jenny Luo Date: Wed, 14 Oct 2026 14:11:19 -0700 Subject: [PATCH 2/3] lora: run the target-name tests in CI and coerce the Hydra list once Review follow-ups on #9. - Move layer selection and the target_modules name checks out of apply_lora_torchtune into _select_lora_targets, which needs no torchtune. The selection and bad-name tests now run in CI (which does not install torchtune); only the LoRALinear and sidecar tests still skip without it. - Convert lora_target_modules from an OmegaConf ListConfig to a plain list in build_training_config, so the config the trainers dump into training_state.json records a list rather than its string form. A bare string passes through unchanged, so apply_lora_torchtune still rejects it. - Update the BioReason demo notebook, which still described lora_target_modules as metadata only. Co-Authored-By: Claude Opus 5.5 --- applications/bioreason/bioreason_demo.ipynb | 2 +- src/train.py | 19 +++++- src/utils/lora_utils.py | 66 ++++++++++++++------- tests/test_lora_target_modules.py | 52 +++++++++------- tests/test_train_coerce.py | 39 +++++++++++- 5 files changed, 127 insertions(+), 51 deletions(-) diff --git a/applications/bioreason/bioreason_demo.ipynb b/applications/bioreason/bioreason_demo.ipynb index 62b56ca..ae7d1ff 100644 --- a/applications/bioreason/bioreason_demo.ipynb +++ b/applications/bioreason/bioreason_demo.ipynb @@ -1969,7 +1969,7 @@ "There's no `projector_type` switch. The DNA projector is always a `ModalityProjector` MLP (fc → GELU → fc → L2 token-norm → + modality embedding), and its **capacity** is set by two knobs: `projector_hidden_mult` (width) and `projector_num_layers` (depth). They're the `W2X`/`D2X` variants in `src/modules/projector.py:VARIANT_MAP`.\n", "\n", "LoRA is applied with torchtune's `apply_lora_torchtune(backbone, rank, alpha)`. That's what the trainer does when `training.lora_enabled=true`.\n", - "> `training.lora_target_modules` is **metadata only**: *every* `nn.Linear` in the backbone except `lm_head` gets an adapter." + "> `training.lora_target_modules` picks the layers that get an adapter: `null` (the default) adapts every `nn.Linear` in the backbone except `lm_head`, and a list of leaf names such as `[q_proj, v_proj]` adapts only those layers." ] }, { diff --git a/src/train.py b/src/train.py index 451175e..5f030c1 100644 --- a/src/train.py +++ b/src/train.py @@ -8,7 +8,7 @@ import numpy as np import torch import torch.distributed as dist -from omegaconf import DictConfig, OmegaConf +from omegaconf import DictConfig, ListConfig, OmegaConf from torch.utils.data import DataLoader # Path Safety @@ -64,6 +64,21 @@ def _coerce_str_or_none(value): return str(value) +def _coerce_listconfig(value): + """Return an OmegaConf ListConfig as a plain list; anything else unchanged. + + Hydra hands list-valued fields over as ListConfig, which json cannot + serialize: the trainers' json.dump(..., default=str) of the config then + records ['q_proj'] in training_state.json as the string "['q_proj']". + Anything else, including None and a bare string, is left for the consumer + to validate (apply_lora_torchtune rejects a bare string; list() here would + split it into characters instead). + """ + if isinstance(value, ListConfig): + return OmegaConf.to_container(value, resolve=True) + return value + + def _is_main_env() -> bool: """Return True when running on rank-0 (or single-process mode).""" rank = int(os.environ.get("RANK", os.environ.get("PMI_RANK", "0"))) @@ -460,7 +475,7 @@ def build_training_config( lora_r=cfg.training.get("lora_r", 32), lora_alpha=cfg.training.get("lora_alpha", 64), lora_dropout=cfg.training.get("lora_dropout", 0.05), - lora_target_modules=cfg.training.get("lora_target_modules", None), + lora_target_modules=_coerce_listconfig(cfg.training.get("lora_target_modules", None)), resume_lora_alpha=cfg.training.get("resume_lora_alpha", None), allow_missing_resume_adapter=cfg.training.get("allow_missing_resume_adapter", False), grpo_num_generations=cfg.training.get("grpo_num_generations", 8), diff --git a/src/utils/lora_utils.py b/src/utils/lora_utils.py index 540cff5..fab6bdb 100644 --- a/src/utils/lora_utils.py +++ b/src/utils/lora_utils.py @@ -25,6 +25,48 @@ def _load_adapter_meta(adapter_path: str) -> dict | None: return json.load(f) +def _select_lora_targets( + backbone: nn.Module, + exclude: tuple[str, ...], + target_modules: list[str] | None, +) -> list[tuple[str, nn.Linear]]: + """Return the (dotted name, nn.Linear) pairs that apply_lora_torchtune replaces. + + Every nn.Linear whose leaf name is not in exclude, restricted to the leaf names + in target_modules when that is not None. Kept free of torchtune so the name + checks can be tested in CI, which does not install it. + + Raises: + TypeError: If target_modules is a bare string rather than a list of names. + ValueError: If target_modules is empty, or names a module that matches no + non-excluded nn.Linear (a typo would otherwise train nothing there). + """ + if isinstance(target_modules, str): + raise TypeError( + f"target_modules must be a list of leaf module names, got the string " + f"{target_modules!r}; write [{target_modules!r}]" + ) + targets = None if target_modules is None else set(target_modules) + if targets is not None and not targets: + raise ValueError("target_modules is empty; use None to adapt every nn.Linear") + + selected = [] + for name, module in backbone.named_modules(): + if isinstance(module, nn.Linear): + leaf = name.split(".")[-1] + if leaf not in exclude and (targets is None or leaf in targets): + selected.append((name, module)) + + if targets is not None: + unmatched = sorted(targets - {name.split(".")[-1] for name, _ in selected}) + if unmatched: + raise ValueError( + f"target_modules {unmatched} matched no nn.Linear in the backbone " + f"(excluded: {list(exclude)})" + ) + return selected + + def apply_lora_torchtune( backbone: nn.Module, rank: int = 32, @@ -60,15 +102,6 @@ def apply_lora_torchtune( """ from torchtune.modules.peft import LoRALinear, get_adapter_params, set_trainable_params - if isinstance(target_modules, str): - raise TypeError( - f"target_modules must be a list of leaf module names, got the string " - f"{target_modules!r}; write [{target_modules!r}]" - ) - targets = None if target_modules is None else set(target_modules) - if targets is not None and not targets: - raise ValueError("target_modules is empty; use None to adapt every nn.Linear") - def _get_parent(model: nn.Module, dotted_name: str): parts = dotted_name.split(".") parent = model @@ -77,20 +110,7 @@ def _get_parent(model: nn.Module, dotted_name: str): return parent, parts[-1] # Collect replacements first to avoid mutating the module tree mid-iteration. - to_replace = [] - for name, module in backbone.named_modules(): - if isinstance(module, nn.Linear): - leaf = name.split(".")[-1] - if leaf not in exclude and (targets is None or leaf in targets): - to_replace.append((name, module)) - - if targets is not None: - unmatched = sorted(targets - {name.split(".")[-1] for name, _ in to_replace}) - if unmatched: - raise ValueError( - f"target_modules {unmatched} matched no nn.Linear in the backbone " - f"(excluded: {list(exclude)})" - ) + to_replace = _select_lora_targets(backbone, exclude, target_modules) replaced = 0 for name, module in to_replace: diff --git a/tests/test_lora_target_modules.py b/tests/test_lora_target_modules.py index f263dfb..238c76c 100644 --- a/tests/test_lora_target_modules.py +++ b/tests/test_lora_target_modules.py @@ -3,7 +3,9 @@ Previously ``training.lora_target_modules`` was written to the adapter sidecar but never reached ``apply_lora_torchtune``, so every nn.Linear except lm_head was adapted regardless; and a Hydra-supplied list crashed ``save_lora_adapter``'s -json.dump. Skipped where torchtune is not installed (CI does not install it). +json.dump. The layer-selection and name-check tests need no torchtune and run in +CI; the LoRALinear and sidecar tests skip where torchtune is not installed (CI +does not install it). """ from __future__ import annotations @@ -12,10 +14,7 @@ import pytest import torch.nn as nn - -pytest.importorskip("torchtune") -from src.utils.lora_utils import apply_lora_torchtune, save_lora_adapter # noqa: E402 -from torchtune.modules.peft import LoRALinear # noqa: E402 +from src.utils.lora_utils import _select_lora_targets, apply_lora_torchtune, save_lora_adapter pytestmark = pytest.mark.unit @@ -35,27 +34,20 @@ def __init__(self): self.lm_head = nn.Linear(8, 16) -def _adapted(model: nn.Module) -> set[str]: - return {name for name, module in model.named_modules() if isinstance(module, LoRALinear)} +_ALL = {f"layers.{i}.{leaf}" for i in range(2) for leaf in ("q_proj", "v_proj", "o_proj")} +_QV = {f"layers.{i}.{leaf}" for i in range(2) for leaf in ("q_proj", "v_proj")} -def test_default_adapts_every_linear_except_lm_head(): - model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0) - assert _adapted(model) == { - f"layers.{i}.{leaf}" for i in range(2) for leaf in ("q_proj", "v_proj", "o_proj") - } +def _selected(target_modules) -> set[str]: + return {name for name, _ in _select_lora_targets(_Tiny(), ("lm_head",), target_modules)} -def test_target_modules_restricts_adaptation_and_trainable_params(): - model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=["q_proj", "v_proj"]) - assert _adapted(model) == {f"layers.{i}.{leaf}" for i in range(2) for leaf in ("q_proj", "v_proj")} - trainable = {name for name, p in model.named_parameters() if p.requires_grad} - assert trainable == { - f"layers.{i}.{leaf}.lora_{ab}.weight" - for i in range(2) - for leaf in ("q_proj", "v_proj") - for ab in ("a", "b") - } +def test_default_selects_every_linear_except_lm_head(): + assert _selected(None) == _ALL + + +def test_target_modules_restricts_selection(): + assert _selected(["q_proj", "v_proj"]) == _QV @pytest.mark.parametrize( @@ -69,10 +61,24 @@ def test_target_modules_restricts_adaptation_and_trainable_params(): ) def test_bad_target_modules_raise(target_modules, error): with pytest.raises(error): - apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=target_modules) + _selected(target_modules) + + +@pytest.mark.parametrize( + ("target_modules", "expected"), [(None, _ALL), (["q_proj", "v_proj"], _QV)] +) +def test_apply_adapts_selected_layers_and_trains_only_lora(target_modules, expected): + LoRALinear = pytest.importorskip("torchtune.modules.peft").LoRALinear + model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=target_modules) + assert { + name for name, module in model.named_modules() if isinstance(module, LoRALinear) + } == expected + trainable = {name for name, p in model.named_parameters() if p.requires_grad} + assert trainable == {f"{name}.lora_{ab}.weight" for name in expected for ab in ("a", "b")} def test_sidecar_records_hydra_list(tmp_path): + pytest.importorskip("torchtune") omegaconf = pytest.importorskip("omegaconf") targets = omegaconf.OmegaConf.create({"t": ["q_proj"]}).t # as cfg.training.get() returns it model = apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=targets) diff --git a/tests/test_train_coerce.py b/tests/test_train_coerce.py index e881dca..46c5012 100644 --- a/tests/test_train_coerce.py +++ b/tests/test_train_coerce.py @@ -1,11 +1,15 @@ """Unit tests for src/train.py helpers. Currently tests _coerce_str_or_none, the defensive coerce that prevents -empty OmegaConf containers from polluting perf.jsonl with "sweep_id": "{}". +empty OmegaConf containers from polluting perf.jsonl with "sweep_id": "{}", +and _coerce_listconfig, which turns a Hydra ListConfig into a plain list so +training_state.json records a list rather than its string form. """ +import json + import pytest -from src.train import _coerce_str_or_none +from src.train import _coerce_listconfig, _coerce_str_or_none def test_none_passes_through(): @@ -52,6 +56,22 @@ def test_integer_stringifies(): assert _coerce_str_or_none(42) == "42" +def test_listconfig_becomes_plain_list(): + """Left as a ListConfig, json.dump(..., default=str) records "['q_proj', 'v_proj']".""" + omegaconf = pytest.importorskip("omegaconf") + # A ListConfig, as cfg.training.get() returns it. + targets = omegaconf.OmegaConf.create({"t": ["q_proj", "v_proj"]}).t + out = _coerce_listconfig(targets) + assert type(out) is list + assert json.dumps(out) == '["q_proj", "v_proj"]' + + +@pytest.mark.parametrize("value", [None, "q_proj", ["q_proj"]]) +def test_non_listconfig_passes_through(value): + """A bare string must reach apply_lora_torchtune intact, which rejects it.""" + assert _coerce_listconfig(value) is value + + # === Hydra wire-through regression tests === # # PR-98 (commit 3ef83a6) added `eval_enabled` to TrainingConfig and @@ -170,3 +190,18 @@ def test_train_py_threads_seed_into_training_config(): "src/train.py missing `seed=seed,` in TrainingConfig(...) — perf.jsonl " "startup_param_count event won't carry the resolved seed" ) + + +def test_train_py_coerces_lora_target_modules(): + """`cfg.training.lora_target_modules` must pass through _coerce_listconfig, + so TrainingConfig holds a plain list rather than an OmegaConf ListConfig. + """ + text = _train_py_text() + pattern = re.compile( + r"lora_target_modules\s*=\s*_coerce_listconfig\(\s*cfg\.training\.get\(\s*['\"]lora_target_modules['\"]", + re.MULTILINE, + ) + assert pattern.search(text), ( + "src/train.py missing `lora_target_modules=_coerce_listconfig(cfg.training.get('lora_target_modules', ...))` " + "— TrainingConfig would hold a ListConfig" + ) From 59064b6560a31647949cc0751476438712b5d41f Mon Sep 17 00:00:00 2001 From: Adarsha Balaji Date: Thu, 15 Oct 2026 20:33:10 +0000 Subject: [PATCH 3/3] lora: list valid target names in the error and log what was adapted When lora_target_modules names a layer that matches nothing, _select_lora_targets now lists the leaf names that can be targeted in the ValueError, e.g. "Available: ['down_proj', 'gate_proj', 'k_proj', ...]". The names come from the backbone actually being adapted, so the list is right for any architecture, and the check is torchtune-free so its test runs in CI. The startup log line in apply_lora_torchtune now breaks the replaced layers down by leaf name, e.g. "replaced 32 linear layers (q_proj: 16, v_proj: 16)", so every run's log shows which layers got adapters. The notebook's recorded output is updated to the new format. Co-Authored-By: Claude Opus 5.5 --- applications/bioreason/bioreason_demo.ipynb | 2 +- src/utils/lora_utils.py | 29 ++++++++++++++------- tests/test_lora_target_modules.py | 21 +++++++++++++++ 3 files changed, 41 insertions(+), 11 deletions(-) diff --git a/applications/bioreason/bioreason_demo.ipynb b/applications/bioreason/bioreason_demo.ipynb index ae7d1ff..f682cd8 100644 --- a/applications/bioreason/bioreason_demo.ipynb +++ b/applications/bioreason/bioreason_demo.ipynb @@ -2033,7 +2033,7 @@ "2026-10-11 23:26:30,490 - src.modules.projector - INFO - Modality embedding: scale=0.0055, expected_norm=0.2489\n", "2026-10-11 23:26:30,491 - src.modules.projector - INFO - ModalityProjector initialized: norm_mode=l2_token, target_norm=0.25, embed_pos=after_norm\n", "2026-10-11 23:26:30,491 - src.model - INFO - Freezing Modality Encoders...\n", - "2026-10-11 23:26:36,815 - src.utils.lora_utils - INFO - torchtune LoRA: replaced 112 linear layers | trainable 12,058,624 / 1,291,845,632 params (0.933%)\n" + "2026-10-11 23:26:36,815 - src.utils.lora_utils - INFO - torchtune LoRA: replaced 112 linear layers (down_proj: 16, gate_proj: 16, k_proj: 16, o_proj: 16, q_proj: 16, up_proj: 16, v_proj: 16) | trainable 12,058,624 / 1,291,845,632 params (0.933%)\n" ] }, { diff --git a/src/utils/lora_utils.py b/src/utils/lora_utils.py index fab6bdb..00aa8a1 100644 --- a/src/utils/lora_utils.py +++ b/src/utils/lora_utils.py @@ -1,6 +1,7 @@ import json import logging import os +from collections import Counter import torch import torch.nn as nn @@ -40,6 +41,7 @@ def _select_lora_targets( TypeError: If target_modules is a bare string rather than a list of names. ValueError: If target_modules is empty, or names a module that matches no non-excluded nn.Linear (a typo would otherwise train nothing there). + The message lists the leaf names that can be targeted. """ if isinstance(target_modules, str): raise TypeError( @@ -50,19 +52,24 @@ def _select_lora_targets( if targets is not None and not targets: raise ValueError("target_modules is empty; use None to adapt every nn.Linear") - selected = [] - for name, module in backbone.named_modules(): - if isinstance(module, nn.Linear): - leaf = name.split(".")[-1] - if leaf not in exclude and (targets is None or leaf in targets): - selected.append((name, module)) + eligible = [ + (name, module) + for name, module in backbone.named_modules() + if isinstance(module, nn.Linear) and name.split(".")[-1] not in exclude + ] + selected = [ + (name, module) + for name, module in eligible + if targets is None or name.split(".")[-1] in targets + ] if targets is not None: - unmatched = sorted(targets - {name.split(".")[-1] for name, _ in selected}) + available = sorted({name.split(".")[-1] for name, _ in eligible}) + unmatched = sorted(targets - set(available)) if unmatched: raise ValueError( - f"target_modules {unmatched} matched no nn.Linear in the backbone " - f"(excluded: {list(exclude)})" + f"target_modules {unmatched} matched no nn.Linear in the backbone. " + f"Available: {available} (excluded: {list(exclude)})" ) return selected @@ -136,8 +143,10 @@ def _get_parent(model: nn.Module, dotted_name: str): n_trainable = sum(p.numel() for p in backbone.parameters() if p.requires_grad) n_total = sum(p.numel() for p in backbone.parameters()) + per_leaf = Counter(name.split(".")[-1] for name, _ in to_replace) logger.info( - f"torchtune LoRA: replaced {replaced} linear layers | " + f"torchtune LoRA: replaced {replaced} linear layers " + f"({', '.join(f'{leaf}: {n}' for leaf, n in sorted(per_leaf.items()))}) | " f"trainable {n_trainable:,} / {n_total:,} params " f"({100.0 * n_trainable / n_total:.3f}%)" ) diff --git a/tests/test_lora_target_modules.py b/tests/test_lora_target_modules.py index 238c76c..2ff94fe 100644 --- a/tests/test_lora_target_modules.py +++ b/tests/test_lora_target_modules.py @@ -11,6 +11,7 @@ from __future__ import annotations import json +import logging import pytest import torch.nn as nn @@ -64,6 +65,12 @@ def test_bad_target_modules_raise(target_modules, error): _selected(target_modules) +@pytest.mark.parametrize("target_modules", [["q_prj"], ["lm_head"]]) +def test_unmatched_error_lists_available_names(target_modules): + with pytest.raises(ValueError, match=r"Available: \['o_proj', 'q_proj', 'v_proj'\]"): + _selected(target_modules) + + @pytest.mark.parametrize( ("target_modules", "expected"), [(None, _ALL), (["q_proj", "v_proj"], _QV)] ) @@ -77,6 +84,20 @@ def test_apply_adapts_selected_layers_and_trains_only_lora(target_modules, expec assert trainable == {f"{name}.lora_{ab}.weight" for name in expected for ab in ("a", "b")} +@pytest.mark.parametrize( + ("target_modules", "summary"), + [ + (None, "replaced 6 linear layers (o_proj: 2, q_proj: 2, v_proj: 2)"), + (["q_proj", "v_proj"], "replaced 4 linear layers (q_proj: 2, v_proj: 2)"), + ], +) +def test_log_reports_adapted_names(caplog, target_modules, summary): + pytest.importorskip("torchtune") + with caplog.at_level(logging.INFO, logger="src.utils.lora_utils"): + apply_lora_torchtune(_Tiny(), rank=2, alpha=4.0, target_modules=target_modules) + assert summary in caplog.text + + def test_sidecar_records_hydra_list(tmp_path): pytest.importorskip("torchtune") omegaconf = pytest.importorskip("omegaconf")