diff --git a/applications/bioreason/bioreason_demo.ipynb b/applications/bioreason/bioreason_demo.ipynb index 62b56ca..f682cd8 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." ] }, { @@ -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/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/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..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 @@ -25,14 +26,66 @@ 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). + The message lists the leaf names that can be targeted. + """ + 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") + + 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: + 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"Available: {available} (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 +96,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 +117,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: @@ -88,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}%)" ) @@ -123,7 +180,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..2ff94fe --- /dev/null +++ b/tests/test_lora_target_modules.py @@ -0,0 +1,109 @@ +"""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 logging + +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", [["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)] +) +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")} + + +@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") + 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" + )