lora: honor lora_target_modules and fix the sidecar save crash

[jenny-luo-0727]

Description

training.lora_target_modules had no effect, and setting it crashed the first checkpoint save.

1. The setting was ignored. TrainingConfig.lora_target_modules was never passed to apply_lora_torchtune, which had no parameter for it. ZoneATrainer and ZoneDTrainer therefore wrapped every nn.Linear except lm_head, whatever the config said. The value was still written to the adapter's .meta.json sidecar, so the sidecar could list target modules that were never adapted.
2. Setting it crashed the first checkpoint save. Hydra passes the list as an OmegaConf ListConfig. save_lora_adapter then raised TypeError: Object of type ListConfig is not JSON serializable when writing the sidecar.

Changes:

• apply_lora_torchtune takes target_modules: list[str] | None = None. None keeps the current behavior: every nn.Linear except exclude. A list restricts LoRA to those leaf names (e.g. [q_proj, v_proj]), matched the same way as the existing exclude.
• A target name that matches no non-excluded nn.Linear raises ValueError, so a typo can't leave layers unadapted without notice. An empty list also raises ValueError. A bare string raises TypeError, since iterating it would split the name into single characters.
• ZoneATrainer and ZoneDTrainer pass config.lora_target_modules through.
• save_lora_adapter converts target_modules to a plain list before json.dump.
• The comments describing the option are updated in config.py, bioreason_sft.yaml and trainer_zone_a.py.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that changes existing behavior)
• Documentation update
• Infrastructure, CI, or tooling

Compatibility

• src.* import paths
• the prism CLI surface
• Hydra config keys
• checkpoint keys
• launcher flags
• None of the above

With lora_target_modules unset (the default, and every shipped config), the result is unchanged. See the default-behaviour check under testing.

How Has This Been Tested?

Run on a NERSC Perlmutter login node, CPU only (torch 2.14.1+cpu, Python 3.12), with requirements/ci.txt, mirroring pr_quality.yml:

ruff check src tests tools scripts applications examples   # ruff 0.16.9: all checks passed
mypy --python-version 3.12 src                              # no issues found in 123 source files
pytest -q -m "not multimodal and not launcher and not integration and not network and not slow and not gpu and not aurora and not perlmutter" tests
# 1354 passed, 29 skipped, 180 deselected, 0 failed

The new tests need torchtune, which requirements/ci.txt does not install, so they skip in the run above. They were run separately with torchtune 0.6.1 (PyPI), torchao 0.5.0 and omegaconf:

pytest tests/test_lora_target_modules.py   # 7 passed; 5 of the 7 fail against the pre-fix lora_utils.py

Default-behaviour check: the old and new apply_lora_torchtune were run on the same seeded toy backbone, with linear layers both with and without bias, and with target_modules unset. All 35 state-dict tensors are bit-identical and the set of trainable parameters is the same.

• Platform tested on: Perlmutter login node, CPU only.
• Tests skipped, and why: the Python 3.10 matrix leg (only 3.12 was run); GPU/Aurora/multimodal/integration-marked tests (deselected by the CI marker expression); no end-to-end BioReason SFT/GRPO run (needs Aurora and the KEGG data).

New Dependencies, Data, or Assets

None.

Checklist

• I have performed a self-review of my own changes
• The code follows the project's style guidelines (ruff check, ruff 0.16.9 per requirements/ci.txt)
• I have commented my code where the intent is not obvious
• I have made corresponding changes to the documentation
• I have added tests that cover this change, or explained why none are needed
• New and existing tests pass locally
• No secrets, credentials, or absolute institutional paths are included

AI/LLM Assistance

• AI/LLM tools were primarily used to generate code or artifacts in this PR
• I have reviewed all such output line-by-line and understand it

The fix and tests were written with Claude Code (Claude Opus 5.5).

Additional Context

The new tests use pytest.importorskip("torchtune"), and requirements/ci.txt does not install torchtune, so CI will skip them. They ran locally against torchtune 0.6.1 (PyPI) with torchao 0.5.0. The sidecar test also needs omegaconf.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GmpnEoNqJpQjR59ie7ZMek

[balajiadarsha]

Thanks, this fixes both bugs as described. I checked the code that uses the adapters and adapting only some layers is safe there: merge_lora_into_base and eval_kegg_aurora.py merge whatever lora pairs are in the saved file, both trainers' watchdogs only check that some LoRA parameter is trainable, and SFT then GRPO with different target lists works because the SFT adapter is merged into the base weights first. I also ran the new tests locally with torchtune 0.6.1 and torchao 0.5.0: 7 passed.

Please address in this PR

1. The BioReason demo notebook now contradicts the code. applications/bioreason/bioreason_demo.ipynb, cell cell-049, still says "training.lora_target_modules is metadata only: every nn.Linear in the backbone except lm_head gets an adapter." After this PR that's wrong, and it's the most user-facing description of the option. Suggest: null adapts every nn.Linear except lm_head, and a list such as [q_proj, v_proj] adapts only those leaf names.
2. CI never runs the new tests (see inline comment on the test file).

Smaller points; fine as follow-ups

3. The ListConfig conversion happens late (see inline comment on save_lora_adapter).
4. The sidecar records target_modules, but nothing checks it on load. load_lora_adapter compares only rank. Now that the option changes which layers get adapters, loading an adapter into a backbone adapted on different layers only logs unexpected keys, and a full SFT resume after changing the list fails with a raw state-dict key mismatch. An expected_target_modules check like the existing expected_rank one, plus a note in the config comment that changing the list makes earlier checkpoints incompatible, would cover it.

Nit: getattr(config, "lora_target_modules", None) could be config.lora_target_modules since the field always exists on TrainingConfig, but it matches the surrounding style, so fine either way.

Handled in a separate PR, #10 (built on this branch, so merge this one first):

• The "matched no nn.Linear" error now lists the leaf names that can be targeted (Available: ['down_proj', 'gate_proj', ...]), taken from the actual backbone, so it's correct for any architecture.
• The startup log line now breaks the replaced layers down by name (replaced 32 linear layers (q_proj: 16, v_proj: 16)).

Items 1 and 2 aren't covered by #10.

[balajiadarsha]

requirements/ci.txt deliberately doesn't install torchtune, so CI skips every test in this file and nothing guards against the bug coming back. The most fragile part, the name checks (typo, empty list, bare string, excluded name) and which layers get selected, doesn't need torchtune. Suggest moving it into a torchtune-free helper in lora_utils.py, e.g. _select_lora_targets(backbone, exclude, target_modules) -> list[tuple[str, nn.Linear]], which apply_lora_torchtune calls. Those tests could then run in CI, and only the LoRALinear and sidecar tests would stay behind importorskip.

[balajiadarsha]

This fixes the crash here, but TrainingConfig.lora_target_modules is still a ListConfig everywhere else at runtime. For example, both trainers write self.config.__dict__ to training_state.json with json.dump(..., default=str), so that file records the string "['q_proj', 'v_proj']" rather than a list. Converting once in src/train.py:463, the way _coerce_str_or_none already handles other fields (e.g. via OmegaConf.to_container), would give every consumer a plain list. Keeping this list() as a fallback for direct callers is fine.

[balajiadarsha]

#10 extends this message to list the available leaf names, which is easier to review there than to add here.

[jenny-luo-0727]

Thank you for the review. Addressed in 5c9e8c2:

1. Notebook: cell-049 now says null adapts every nn.Linear except lm_head, and a list such as [q_proj, v_proj] adapts only those leaf names.
2. CI coverage: layer selection and the name checks moved into a torchtune-free _select_lora_targets, which apply_lora_torchtune calls. The selection and bad-name tests now run in CI. Only the LoRALinear and sidecar tests stay behind importorskip. I checked that reverting the fix makes a CI-runnable test fail.
3. ListConfig: I picked this one up here too. build_training_config now converts lora_target_modules to a plain list via OmegaConf.to_container, so training_state.json records a real list. A bare string passes through unchanged so apply_lora_torchtune still raises its TypeError. The list() in save_lora_adapter stays as the fallback. Tests are in tests/test_train_coerce.py.

I've left 4 (expected_target_modules on load) for a follow-up, as you suggested.

Heads-up for #10: moving the selection loop into _select_lora_targets means #10's eligible / Available: change will need to land in the helper when you rebase. Sorry for the churn!