[BIO-303] support LoRA in evo2 mbridge

[balvisio]

Description

Usage

TODO: Add code snippet

Type of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Refactor
• Documentation update
• Other (please describe):

CI Pipeline Configuration

Configure CI behavior by applying the relevant labels. By default, only basic unit tests are run.

• ciflow:skip - Skip all CI tests for this PR
• ciflow:notebooks - Run Jupyter notebooks execution tests
• ciflow:slow - Run slow single GPU integration tests marked as @pytest.mark.slow
• ciflow:all - Run all tests (unit tests, slow tests, and notebooks). This label can be used to enforce running all framework tests.
• ciflow:all-recipes - Run tests for all recipes (under bionemo-recipes). This label can be used to enforce running tests for all recipes.

Unit tests marked as @pytest.mark.multi_gpu or @pytest.mark.distributed are not run in the PR pipeline.

For more details, see CONTRIBUTING

Note

By default, only basic unit tests are run. Add appropriate labels to enable an additional test coverage.

Authorizing CI Runs

We use copy-pr-bot to manage authorization of CI
runs on NVIDIA's compute resources.

• If a pull request is opened by a trusted user and contains only trusted changes, the pull request's code will
  automatically be copied to a pull-request/ prefixed branch in the source repository (e.g. pull-request/123)
• If a pull request is opened by an untrusted user or contains untrusted changes, an NVIDIA org member must leave an
  /ok to test comment on the pull request to trigger CI. This will need to be done for each new commit.

Triggering Code Rabbit AI Review

To trigger a code review from code rabbit, comment on a pull request with one of these commands:

• @coderabbitai review - Triggers a standard review
• @coderabbitai full review - Triggers a comprehensive review

See https://docs.coderabbit.ai/reference/review-commands for a full list of commands.

Pre-submit Checklist

• I have tested these changes locally
• I have updated the documentation accordingly
• I have added/updated tests as needed
• All existing tests pass successfully

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: d3e73b83-54be-4087-b21a-f8ba9ded6da8

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch dev/ba/BIO-303-support-lora-in-evo2-mbridge

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[jstjohn]

Looks good but the biggest gap I see is around documenting how to ship and ingest a LORA checkpoint. Is it separate from the rest of the model or is it a combined checkpoint with everything in it including the original weights? I'd say call this out in the docs etc so people know what to expect. The other gap I see is the head swap use case which is pretty common for fine tuning. For now you could mention theoretically how you would do it. I think it would end up being a new model type or a new training script that has the desired structure in code, then I'm not 100% how to load the existing weights into it while handling the missing new head weight or changed shape gracefully, maybe it's writing a new checkpoint converter that does the necessary weight mugging? Anyways calling out in the readme how that would be done will probably save us support time long term.

Thanks!

[balvisio]

@jstjohn I added the documentation on how LoRA adapters are saved and how to use them along with the base checkpoint for inference. For the modification of the model head I will add as part of the next PR that contains a LoRA focused Jupyter notebook.