Skip to content

lora: list valid target names in the error and log what was adapted (follow-up to #9) - #10

Open
balajiadarsha wants to merge 3 commits into
mainfrom
lora-target-diagnostics
Open

balajiadarsha wants to merge 3 commits into
mainfrom
lora-target-diagnostics

Conversation

@balajiadarsha

@balajiadarsha balajiadarsha commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

Description

Follow-up to #9, which makes training.lora_target_modules take effect. This branch is built on #9's latest commit (5c9e8c2), so merge #9 first. Until then the diff below also shows #9's changes. The only new commit is the last one (59064b6).

#9 checks target names and raises ValueError when a name matches nothing, but the error doesn't say which names would work, and the run log doesn't say which layers actually got adapters. This PR adds both:

  1. List valid names in the error (in _select_lora_targets, the torchtune-free helper from lora: honor lora_target_modules and fix the sidecar save crash #9). When a target matches no non-excluded nn.Linear, the message now lists the leaf names that can be targeted:
    ValueError: target_modules ['q_prj'] matched no nn.Linear in the backbone. Available: ['down_proj', 'gate_proj', 'k_proj', 'o_proj', 'q_proj', 'up_proj', 'v_proj'] (excluded: ['lm_head'])
    
    The list comes from the backbone being adapted, so it's correct for whatever architecture is loaded.
  2. Show what was adapted in the log (in apply_lora_torchtune). The existing startup line now breaks the replaced layers down by leaf name:
    torchtune LoRA: replaced 32 linear layers (q_proj: 16, v_proj: 16) | trainable ... params (...%)
    

Which layers get selected is unchanged: the selection loop in _select_lora_targets is split into a list of eligible layers and the filtered list built from it, giving the same result. The notebook's recorded apply_lora_torchtune output (bioreason_demo.ipynb, section 9b) is updated to the new log format.

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

Only the text of one error message and one INFO log line changes.

How Has This Been Tested?

Aurora login node, CPU only. Rerun after rebasing onto #9's 5c9e8c2:

# without torchtune (as in CI)
pytest tests/test_lora_target_modules.py tests/test_train_coerce.py   # 27 passed, 5 skipped (torchtune)
# with torchtune 0.6.1 (frameworks 2026.1 venv)
pytest tests/test_lora_target_modules.py tests/test_train_coerce.py   # 32 passed
# with only this commit's lora_utils.py change reverted: the 4 new tests fail
ruff check src/utils/lora_utils.py tests/test_lora_target_modules.py  # all checks passed

New tests:

  • test_unmatched_error_lists_available_names, for a typo and for the excluded lm_head. It calls _select_lora_targets, so it runs in CI.
  • test_log_reports_adapted_names, for the default (all layers) and for a restricted list. It needs torchtune, so CI skips it.

End-to-end: #9's 5c9e8c2 was run end to end on Aurora (SFT with [q_proj, v_proj], GRPO with and without targets, eval, resume) and passed. This commit only changes an error message and a log line, so it has not had its own training run yet.

  • Platform tested on: Aurora login node, CPU only.
  • Tests skipped, and why: test_log_reports_adapted_names needs torchtune, which requirements/ci.txt doesn't install.

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

Written with Claude Code (Claude Opus 5.5).

🤖 Generated with Claude Code

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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01GmpnEoNqJpQjR59ie7ZMek
jenny-luo-0727 and others added 2 commits October 14, 2026 14:11
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 <noreply@anthropic.com>
When lora_target_modules names a layer that matches nothing,
_select_lora_targets now lists the leaf names that can be targeted in
the ValueError, e.g. "Available: ['down_proj', 'gate_proj', 'k_proj',
...]". The names come from the backbone actually being adapted, so the
list is right for any architecture, and the check is torchtune-free so
its test runs in CI.

The startup log line in apply_lora_torchtune now breaks the replaced
layers down by leaf name, e.g. "replaced 32 linear layers (q_proj: 16,
v_proj: 16)", so every run's log shows which layers got adapters. The
notebook's recorded output is updated to the new format.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants