diff --git a/applications/bioreason/bioreason_demo.ipynb b/applications/bioreason/bioreason_demo.ipynb index 62b56ca..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` 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` 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 a439b8d..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 = auto-detect all nn.Linear layers (excluding lm_head) +# 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/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/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/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..fab6bdb 100644 --- a/src/utils/lora_utils.py +++ b/src/utils/lora_utils.py @@ -25,14 +25,60 @@ 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, 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,9 +89,16 @@ 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 @@ -57,12 +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: - to_replace.append((name, module)) + to_replace = _select_lora_targets(backbone, exclude, target_modules) replaced = 0 for name, module in to_replace: @@ -123,7 +171,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_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/tests/test_lora_target_modules.py b/tests/test_lora_target_modules.py new file mode 100644 index 0000000..238c76c --- /dev/null +++ b/tests/test_lora_target_modules.py @@ -0,0 +1,88 @@ +"""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. 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 + +import json + +import pytest +import torch.nn as nn +from src.utils.lora_utils import _select_lora_targets, apply_lora_torchtune, save_lora_adapter + +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) + + +_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 _selected(target_modules) -> set[str]: + return {name for name, _ in _select_lora_targets(_Tiny(), ("lm_head",), target_modules)} + + +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( + ("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): + _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) + 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"]} 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" + ) 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()