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 b427d6d9af6f4f33d88ada6bb0b3f6c9eaaeb6fe Mon Sep 17 00:00:00 2001 From: Adarsha Balaji Date: Thu, 15 Oct 2026 20:45:39 +0000 Subject: [PATCH 3/3] bioreason: default GRPO LoRA to q_proj/v_proj and add a target lister bioreason_grpo.yaml now ships lora_target_modules: [q_proj, v_proj], so GRPO adapts only the attention query/value projections (OLMo-1B: 32 of 112 layers) instead of every nn.Linear. bioreason_sft.yaml declares the key as null, keeping SFT's behavior. Declaring it in both lets training.lora_target_modules=[...] be overridden without Hydra's "+". tools/list_lora_targets.py prints the leaf names that can be targeted for any model, with layer counts. It builds the backbone on the meta device (config.json only, no weights) and uses the trainer's own _select_lora_targets. Both YAMLs and the notebook point to it. Co-Authored-By: Claude Opus 5.5 --- applications/bioreason/bioreason_demo.ipynb | 2 +- src/conf/training/bioreason_grpo.yaml | 8 +++ src/conf/training/bioreason_sft.yaml | 6 +- tests/test_bioreason_lora_targets_config.py | 64 ++++++++++++++++++ tools/list_lora_targets.py | 75 +++++++++++++++++++++ 5 files changed, 153 insertions(+), 2 deletions(-) create mode 100644 tests/test_bioreason_lora_targets_config.py create mode 100644 tools/list_lora_targets.py diff --git a/applications/bioreason/bioreason_demo.ipynb b/applications/bioreason/bioreason_demo.ipynb index ae7d1ff..a86b145 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` 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." + "> `training.lora_target_modules` picks the layers that get an adapter: `null` 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. `bioreason_sft.yaml` ships `null`; `bioreason_grpo.yaml` ships `[q_proj, v_proj]`. The names differ between architectures; `python tools/list_lora_targets.py model=prism_olmo_1b` lists the valid ones for a model without loading its weights." ] }, { diff --git a/src/conf/training/bioreason_grpo.yaml b/src/conf/training/bioreason_grpo.yaml index d4306ff..11bdf82 100644 --- a/src/conf/training/bioreason_grpo.yaml +++ b/src/conf/training/bioreason_grpo.yaml @@ -34,6 +34,14 @@ lora_enabled: true lora_r: 16 lora_alpha: 32 lora_dropout: 0.0 +# Leaf names of the backbone nn.Linear layers that get a GRPO adapter. +# [q_proj, v_proj] adapts only the attention query/value projections +# (OLMo-1B: 32 of 112 layers). null adapts every nn.Linear except lm_head. +# Names differ between architectures; list the valid ones for a model with: +# python tools/list_lora_targets.py model=prism_olmo_1b +# A misspelled name fails at startup. Independent of the SFT adapter's +# targets: that adapter is merged into the base weights first. +lora_target_modules: [q_proj, v_proj] # alpha the resume_weights_only checkpoint's lora_adapter.pt was trained with # (see bioreason_sft.yaml lora_alpha). Used only to merge that adapter into the diff --git a/src/conf/training/bioreason_sft.yaml b/src/conf/training/bioreason_sft.yaml index 4a33c47..d8dcca8 100644 --- a/src/conf/training/bioreason_sft.yaml +++ b/src/conf/training/bioreason_sft.yaml @@ -72,7 +72,11 @@ lora_enabled: true lora_r: 32 lora_alpha: 64 lora_dropout: 0.05 -# lora_target_modules: null # null = all nn.Linear layers (excluding lm_head); or leaf names, e.g. [q_proj, v_proj] +# Leaf names of the backbone nn.Linear layers that get an adapter. +# null adapts every nn.Linear except lm_head; a list such as [q_proj, v_proj] +# adapts only those. List the valid names for a model with: +# python tools/list_lora_targets.py model=prism_olmo_1b +lora_target_modules: null # Data format: include reasoning traces in Stage 2 so LoRA learns chain-of-thought. use_reasoning_traces: true diff --git a/tests/test_bioreason_lora_targets_config.py b/tests/test_bioreason_lora_targets_config.py new file mode 100644 index 0000000..8c61726 --- /dev/null +++ b/tests/test_bioreason_lora_targets_config.py @@ -0,0 +1,64 @@ +"""The BioReason training YAMLs declare lora_target_modules, and the listing tool works. + +Declaring the key lets ``training.lora_target_modules=[...]`` override it +without Hydra's ``+`` prefix (struct mode rejects undeclared keys). GRPO ships +``[q_proj, v_proj]``; SFT ships ``null`` (every nn.Linear except lm_head). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch.nn as nn +from hydra import compose, initialize_config_dir +from src.train import _coerce_listconfig +from tools.list_lora_targets import count_targets + +pytestmark = pytest.mark.unit + +CONF = Path(__file__).resolve().parents[1] / "src" / "conf" + + +def _targets(*overrides: str): + with initialize_config_dir(version_base=None, config_dir=str(CONF)): + cfg = compose(config_name="config", overrides=["model=prism_olmo_1b", *overrides]) + return _coerce_listconfig(cfg.training.lora_target_modules) + + +@pytest.mark.parametrize( + ("training", "expected"), + [("bioreason_grpo", ["q_proj", "v_proj"]), ("bioreason_sft", None)], +) +def test_shipped_default(training, expected): + assert _targets(f"training={training}") == expected + + +@pytest.mark.parametrize("training", ["bioreason_grpo", "bioreason_sft"]) +@pytest.mark.parametrize( + ("override", "expected"), + [("[k_proj,o_proj]", ["k_proj", "o_proj"]), ("null", None)], +) +def test_plain_override_needs_no_plus(training, override, expected): + targets = _targets(f"training={training}", f"training.lora_target_modules={override}") + assert targets == expected + assert targets is None or type(targets) is list + + +class _Block(nn.Module): + def __init__(self): + super().__init__() + self.q_proj = nn.Linear(8, 8) + self.v_proj = nn.Linear(8, 8) + self.act = nn.GELU() + + +class _Tiny(nn.Module): + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Block(), _Block(), _Block()]) + self.lm_head = nn.Linear(8, 16) + + +def test_count_targets_lists_linear_leaf_names_except_lm_head(): + assert count_targets(_Tiny()) == {"q_proj": 3, "v_proj": 3} diff --git a/tools/list_lora_targets.py b/tools/list_lora_targets.py new file mode 100644 index 0000000..60afbca --- /dev/null +++ b/tools/list_lora_targets.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python3 +"""List the layer names ``training.lora_target_modules`` accepts for a model. + +LoRA targets are leaf names of the backbone's ``nn.Linear`` layers (e.g. +``q_proj``), and they differ between architectures. This builds the backbone +on the meta device (only its config.json is read, no weights) and prints each +name that apply_lora_torchtune would accept, with how many layers carry it. +The selection is the same ``_select_lora_targets`` the trainer uses. + + python tools/list_lora_targets.py model=prism_olmo_1b + python tools/list_lora_targets.py --backbone-id Qwen/Qwen3-1.7B +""" + +from __future__ import annotations + +import argparse +import sys +from collections import Counter +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +import torch.nn as nn +from src.utils.lora_utils import _select_lora_targets + +EXCLUDE = ("lm_head",) # apply_lora_torchtune's default + + +def count_targets(backbone: nn.Module, exclude: tuple[str, ...] = EXCLUDE) -> Counter[str]: + """Leaf name -> number of nn.Linear layers that LoRA would adapt by default.""" + return Counter(name.split(".")[-1] for name, _ in _select_lora_targets(backbone, exclude, None)) + + +def _backbone_id_from_hydra(overrides: list[str]) -> str: + from hydra import compose, initialize_config_dir + + with initialize_config_dir(config_dir=str(ROOT / "src" / "conf"), version_base="1.2"): + cfg = compose(config_name="config", overrides=overrides) + return cfg.model.backbone_id + + +def _build_on_meta(backbone_id: str) -> nn.Module: + import torch + from transformers import AutoConfig, AutoModelForCausalLM + + config = AutoConfig.from_pretrained(backbone_id, trust_remote_code=True) + with torch.device("meta"): + return AutoModelForCausalLM.from_config(config, trust_remote_code=True) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "overrides", nargs="*", help="Hydra overrides selecting the model, e.g. model=prism_olmo_1b" + ) + parser.add_argument("--backbone-id", help="HF model id or local path; skips Hydra") + args = parser.parse_args() + + backbone_id = args.backbone_id or _backbone_id_from_hydra(args.overrides) + counts = count_targets(_build_on_meta(backbone_id)) + + print(f"Backbone: {backbone_id}") + print(f"LoRA-targetable nn.Linear leaf names (excluded: {list(EXCLUDE)}):") + for leaf, n in sorted(counts.items()): + print(f" {leaf:<24} {n:>4} layers") + print(f" {'total':<24} {sum(counts.values()):>4} layers") + print("\nExample: training.lora_target_modules=[q_proj,v_proj] (null adapts all of the above)") + + +if __name__ == "__main__": + main()