bioreason: default GRPO LoRA to q_proj/v_proj and add a target lister

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."
    ]
   },
   {

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

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

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

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),

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

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:

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)

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}

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"]}