Add support on num steps for learning rate scheduler

[sichu2023]

Decouple num_steps with number of steps in learning rate scheduler.

python sub-packages/bionemo-esm2/src/bionemo/esm2/scripts/train_esm2.py ... --learning-rate-num-steps=...

Or taking the pydantic route

export MY_DATA_SOURCE="ngc"

# The fastest transformer engine environment variables in testing were the following two
TEST_DATA_DIR=$(download_bionemo_data esm2/testdata_esm2_pretrain:2.0 --source $MY_DATA_SOURCE); \
bionemo-esm2-recipe \
--train-cluster-path ${TEST_DATA_DIR}/2024_03_sanity/train_clusters_sanity.parquet     \
--train-database-path ${TEST_DATA_DIR}/2024_03_sanity/train_sanity.db     \
--valid-cluster-path ${TEST_DATA_DIR}/2024_03_sanity/valid_clusters.parquet     \
--valid-database-path ${TEST_DATA_DIR}/2024_03_sanity/validation.db     \
--result-dir ./results     \
--dest my_config.json \
--recipe 8m \
--scheduler-max-steps=500000 \
--max-steps=10000

bionemo-esm2-train \
--data-config-t bionemo.esm2.run.config_models.ESM2DataConfig \
--model-config-t bionemo.esm2.run.config_models.ExposedESM2PretrainConfig \
--config my_config.json

[sichu2023]

/build-ci

[sichu2023]

@skothenhill-nv Do you know where I should update TrainingConfig to mirror the change to train.py?
https://github.com/NVIDIA/bionemo-framework/blob/sichu/esm2_learning_rate_num_steps/sub-packages/bionemo-llm/src/bionemo/llm/train.py#L216

[sichu2023]

/build-ci

[skothenhill-nv]

I didnt quite figure out the right abstraction for schedulers + optimizers. When you do this, if there is an approach that stands out to you please let me know! To answer your question:

@skothenhill-nv Do you know where I should update TrainingConfig to mirror the change to train.py? https://github.com/NVIDIA/bionemo-framework/blob/sichu/esm2_learning_rate_num_steps/sub-packages/bionemo-llm/src/bionemo/llm/train.py#L216

Schedulers are defined in the train method, so you'll want to add its usage here:

1. https://github.com/NVIDIA/bionemo-framework/blob/sichu/esm2_learning_rate_num_steps/sub-packages/bionemo-llm/src/bionemo/llm/train.py#L204

2. TrainingConfig is defined here: https://github.com/NVIDIA/bionemo-framework/blob/sichu/esm2_learning_rate_num_steps/sub-packages/bionemo-llm/src/bionemo/llm/run/config_models.py#L262

add it as a field and update the recipes to do the right thing.

[sichu2023]

/build-ci

[sichu2023]

/build-ci

[sichu2023]

/build-ci

[sichu2023]

/build-ci

[sichu2023]

/build-ci

[sichu2023]

/build-ci

[sichu2023]

/build-ci

[sichu2023]

/build-ci