unify the implementation of early training termination across BioNeMo subpackages and update benchmarks

[dorotat-nv]

Description

Current implementations (Evo2 and ESM2) use different approaches to stop training at specific steps while maintaining the full learning rate schedule or other characteristics. Trying to unify it

Evo2: Uses checkpoint mechanism to stop training after K steps
ESM2: Implements a different solution in train_esm2.py

Evo2: https://github.com/NVIDIA/bionemo-framework/blob/a643d8c073262a542c7b94cc39648173d[…]87c1d8e/sub-packages/bionemo-evo2/src/bionemo/evo2/run/train.py

ESM2 https://github.com/NVIDIA/bionemo-framework/blob/a643d8c073262a542c7b94cc39648173d[…]ub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py

Addressing issue: #749

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:

• SKIP_CI - Skip all continuous integration tests
• INCLUDE_NOTEBOOKS_TESTS - Execute notebook validation tests in pytest
• INCLUDE_SLOW_TESTS - Execute tests labelled as slow in pytest for extensive testing

Note

By default, the notebooks validation tests are skipped unless explicitly enabled.

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.

Usage

TODO: Add code snippet

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

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[sichu2023]

LGTM

[dorotat-nv]

/ok to test

[dorotat-nv]

/ok to test

[codecov-commenter]

Codecov Report

❌ Patch coverage is 50.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 84.34%. Comparing base (f320920) to head (a25f6a1).
⚠️ Report is 441 commits behind head on main.

Files with missing lines	Patch %	Lines
...ackages/bionemo-evo2/src/bionemo/evo2/run/train.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #803   +/-   ##
=======================================
  Coverage   84.33%   84.34%           
=======================================
  Files         137      137           
  Lines        8627     8626    -1     
=======================================
  Hits         7276     7276           
+ Misses       1351     1350    -1     

Files with missing lines	Coverage Δ
...ionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py	93.79% <100.00%> (+0.04%)	⬆️
...ackages/bionemo-evo2/src/bionemo/evo2/run/train.py	16.66% <0.00%> (-0.37%)	⬇️

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[dorotat-nv]

/ok to test