Allow finetuning ESM2 with [un]frozen encoder

[farhadrgh]

Description

The overlap_grad_reduce=True causes communication issue in gradient synchronization step when the encoder parameters are not frozen.

AssertionError: Communication call has not been issued for this bucket (79/84 params have grad available)

This PR changes the default for overlap_grad_reduce and exposes config.encoder_frozen to optionally [un]freeze the encoder parameters.

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

Note

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

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

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 86.63%. Comparing base (eb89056) to head (af1619e).

✅ All tests successful. No failed tests found.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #620   +/-   ##
=======================================
  Coverage   86.63%   86.63%           
=======================================
  Files         116      116           
  Lines        6966     6967    +1     
=======================================
+ Hits         6035     6036    +1     
  Misses        931      931           

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

[sichu2023]

LGTM