Add federated BioNeMo Evo2 fine-tuning example

[holgerroth]

What this adds

This adds a compact BioNeMo example for federated LoRA fine-tuning of evo2/1b-8k-bf16:1.0 on three-class splice-site classification.

The maintained workflow:

1. Prepares deterministic simulated-site partitions from the pinned revised Nucleotide Transformer benchmark.
2. Converts and content-binds the Evo2 checkpoint for Megatron Bridge.
3. Creates one common LoRA and classification-head initialization.
4. Runs sample-weighted FedAvg with fresh external BioNeMo trainer processes scheduled sequentially on one GPU.
5. Reloads the final global checkpoint and evaluates all 3,000 official test rows.

Only the LoRA adapters and classification head cross the federation boundary as CPU FP32 tensors. The BF16 backbone remains frozen at each site. The example validates tensor schemas and provenance, synchronizes optimizer master parameters after loading global weights, rejects no-op runs, checks the frozen backbone, audits data leakage, and requires exact evaluation coverage.

All implementation changes are contained under examples/advanced/bionemo/evo2 and its tests; no shared NVFlare API changes are required. See the README and walkthrough.

Maintenance scope

Review feedback correctly identified that the original version included too much experiment infrastructure for a maintained example. This revision removes persistent optimizer and RNG state, continuation and checkpoint chains, checkpoint averaging, local-only and head-only campaign modes, model selection, baseline summarization, and benchmark-report generation.

The scope cut reduces the branch by more than 5,000 net lines.

Each federated task now reconstructs Evo2, loads the current global trainable tensors, creates fresh local training state, returns its update, and exits. The maintained lesson is the BioNeMo/Megatron integration and federated adapter exchange. Longer benchmark reproduction remains research tooling outside this example.

The environment is pinned to specific BioNeMo Recipes, Megatron Bridge, causal-conv1d, dataset, and Evo2 checkpoint revisions. Updating those dependencies should include a GPU smoke test of the real training path.

Precursor H100 evidence

Before the scope reduction, a precursor snapshot ran a longer three-site IID campaign and matched independent local-only training budgets. These measurements motivate the example, but the corresponding averaging, baseline, and reporting machinery is no longer part of the maintained workflow.

Endpoint	Validation accuracy / macro-F1	Test accuracy / macro-F1
Raw FL round 36	0.9607 / 0.9606	0.9533 / 0.9535
FL mean, rounds 27–36	0.9617 / 0.9617	0.9560 / 0.9562
Local site 1	0.9163 / 0.9165	0.9080 / 0.9083
Local site 2	0.9110 / 0.9117	0.9010 / 0.9023
Local site 3	0.8680 / 0.8690	0.8497 / 0.8514
Descriptive local mean	0.8984 / 0.8991	0.8862 / 0.8873

The raw final FL model exceeded the validation-selected best local model on test by 4.53 percentage points in accuracy and 4.52 points in macro-F1. It exceeded the descriptive local mean by 6.71 and 6.62 points, respectively. The averaged FL endpoint was 1.03 accuracy points below the BioNeMo tutorial's reported 0.9663 result.

Fixed precursor checkpoint identities:

• Raw FL round 36: db62e79e96edcbd1aebdd97bbe31b36100b05f442d8c9ca83f84883764ae787c
• Validation-selected uniform mean of rounds 27–36: 726ce82310b453b81ea1cd0d08d3bb37b839c3f74cfdcfa301d81a1acd0dad39

The averaged endpoint was selected using validation results. The local mean is descriptive and is neither an ensemble nor a model trained on combined local data. These results came from simulated IID partitions on a precursor implementation, and official test results were observed during development. They demonstrate feasibility and collaboration benefit in this controlled experiment; they do not establish broad gains across institutions or heterogeneous genomic datasets.

Validation

• 103 focused Evo2 tests passed.
• 2 existing cross-example NeMo PEFT tests passed, for 105 relevant tests total.
• A two-client, one-round external-process simulator run passed with the mock BioNeMo backend.
• Manifest-bound mock evaluation passed.
• Repository style check, license check, notebook validation, and git diff --check passed.
• The notebook contains eight parseable code cells and no stored outputs.
• The real H100 training path passed at 00e3d4506efb8e12f0b0a7ad9c37c56700d69166: two BioNeMo clients completed one round with four local steps; 199 of 206 trainable tensors changed; both clients verified the frozen backbone; the final checkpoint reloaded and evaluated all 3,000 official test rows; and the GPU had no process or allocated memory before or after the run. The current head differs only in test timeout handling.

The four-step smoke produced 0.3447 accuracy and 0.2366 macro-F1. This deliberately short run validates the integration path, tensor exchange, checkpoint reload, and teardown; it is not a quality benchmark. One warmup step had produced no BF16-visible update, so the example now documents four steps and rejects a completed run whose global trainable state is unchanged.

The final H100 smoke reused a known-good previously converted Evo2 base checkpoint. Its complete directory content was bound to initialization, training, and evaluation metadata, but it predates the current nvflare_conversion_provenance.json format. The smoke therefore validates the final training and evaluation path; checkpoint reconversion was not rerun.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Quick tests passed locally by running ./runtest.sh.
• In-line docstrings updated.
• Documentation updated.

[holgerroth]

/build

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge; no outstanding correctness, security, or repository-rule issue was identified.

Summary

This PR adds a maintained federated BioNeMo Evo2 LoRA fine-tuning example for splice-site classification.

• Prepares deterministic, manifest-bound site partitions with leakage auditing.
• Converts and validates a pinned Evo2 checkpoint and exports a common LoRA/head initialization.
• Runs sequential one-GPU, sample-weighted FedAvg using fresh external trainer processes.
• Strictly validates exchanged tensors, frozen-backbone behavior, final model changes, provenance, and exact test-set coverage.
• Includes focused tests, pinned container inputs, documentation, and an executable walkthrough.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    D[Prepare pinned dataset] --> M[Manifest and leakage audit]
    B[Convert pinned Evo2 checkpoint] --> P[Validate payload and provenance]
    M --> I[Create common LoRA and head initialization]
    P --> I
    I --> S[NVFlare server]
    S --> C1[Site 1 fresh trainer]
    S --> C2[Site 2 fresh trainer]
    S --> C3[Site 3 fresh trainer]
    C1 -->|CPU FP32 DIFF| A[Exact-schema sample-weighted FedAvg]
    C2 -->|CPU FP32 DIFF| A
    C3 -->|CPU FP32 DIFF| A
    A --> S
    S --> G[Final global trainable checkpoint]
    G --> E[Reload and evaluate all official test rows]

Loading

_{Reviews (2) · Last reviewed commit: "Merge branch 'main' into feat/evo2-bione..."}

[holgerroth]

Reviewed c130a0f. Found one notebook configuration issue, detailed inline. Validation: all 150 focused Evo2 unit tests passed; a two-client, two-round CPU external-process smoke run and mock evaluation passed, with sample-weighted aggregation verified. GPU training and the Docker build were not rerun.

[holgerroth]

/build

[holgerroth]

Re-reviewed ee4adfd. The previous notebook initialization finding is fixed: non-default seed, sequence length, and LoRA settings now reach initialization and evaluation consistently. All 169 focused Evo2 tests passed. A fresh two-client/two-round CPU external-process smoke run, mock evaluation, and a continuation at logical round 2 passed; weighted tensor values and the preserved optimizer signature were verified. One remaining notebook reporting issue is noted inline. GPU training and the Docker build were not rerun.

[holgerroth]

/build

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 67.73%. Comparing base (5f09f3b) to head (eb5cfdc).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5289      +/-   ##
==========================================
+ Coverage   67.72%   67.73%   +0.01%     
==========================================
  Files        1023     1023              
  Lines      106444   106444              
==========================================
+ Hits        72089    72101      +12     
+ Misses      34355    34343      -12     

Flag	Coverage Δ
unit-tests	67.73% <ø> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[holgerroth]

/build

[holgerroth]

/build

[holgerroth]

/build

[holgerroth]

/build