fixing the ESM2 checkpointing issue

[polinabinder1]

Description

This addresses: #757

Type of changes

In the original code the optimizer was not saved in the checkpoint, but it is expected in the megatron strategy. Saving the optimizer is added to the checkpoint callback and the test has been updated. Changes were made to have the checkpointing callback in the nemo logger. This is the same as the training path in sub-packages/bionemo-llm/src/bionemo/llm/train.py.

[trvachov]

@jwilber can you review as well please?

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.39%. Comparing base (3936231) to head (54e9a8e).
Report is 65 commits behind head on main.

✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #842      +/-   ##
==========================================
+ Coverage   84.37%   84.39%   +0.02%     
==========================================
  Files         138      138              
  Lines        8690     8690              
==========================================
+ Hits         7332     7334       +2     
+ Misses       1358     1356       -2     

Files with missing lines	Coverage Δ
...ionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py	93.79% <ø> (ø)

... and 1 file with indirect coverage changes

[skothenhill-nv]

There is similar logic here: https://github.com/NVIDIA/bionemo-framework/blob/main/sub-packages/bionemo-llm/src/bionemo/llm/train.py#L59

for the pydantic based workflow. but checkpoints in this code-path are passed into setup_nemo_lightning_logger but without the additional arguments to the callback- seems like we do not see the same error.

Do you understand what's going on here? What are the tradeoffs when passing the ModelCheckpoint callback via Trainer vs setup_nemo_lightning_logger. It makes me uncomfortable that we have hit a divergent point.

[skothenhill-nv]

Approved- but please address the divergence with sub-packages/bionemo-llm/src/bionemo/llm/train.py .

Ideally all of our models would use the same workflow. We should understand the difference and if its a reason we saw this bug.

Or maybe we should document in the description of this PR the differences two code paths?

[trvachov]

Nit: can you fix the type in the PR title?

[jwilber]

Added style comment, otherwise LGTM