[BUG] Evo2 finetune with LoRA does not work

[dorotat-nv]

BioNeMo Framework Version

0d162d5

Bug Description

When running Evo2 finetune with LoRA, the metrics (e.g., val_loss, other validation metrics, and reduced_train_loss) do not decrease as expected and instead remain almost fixed across training steps. The correct behavior is observed for Evo2 finetunes without LoRA - ie the metrics decrease. The Evo2 finetunes both with and without LoRA are defined in link

The PR with the external contribution that was merged https://github.com/NVIDIA/bionemo-framework/pull/980 

1. Did not include any unit tests to verify LoRA functionality for Evo2, even though it was shown to work on a single GPU (single node).
2. Was not validated against finetuning without LoRA — not even by the author. There is no supporting evidence that it runs correctly or produces meaningful results, making it highly experimental.

Tasks to do

1. Create a unit test demonstrating that LoRA finetuning is failing.

2. Fix the LoRA finetuning functionality.

3. Ensure the unit test passes once the fix is in place.

4. If the issue cannot be resolved immediately, add a notification to the training script indicating that LoRA training is currently broken.

Steps to Reproduce

1. Build a container for 0d162d5ffe97eee471d259e13d82f212f5b4dc5d

2. Set up job (1 node) on eos and mount /data/:/lustre/fsw/healthcareeng_bionemo/jet/data

3. Run the finetuning command for Evo2 for LoRA

train_evo2 -d /data/evo2/training_data_config.yaml \ --dataset-dir=/data/evo2/preprocessed_data \ --ckpt-dir=/data/evo2/checkpoints/nemo2_evo2_1b_8k \ --lora-finetune \ --model-size=1b \ --max-steps=1000 \ --experiment-name=evo2-lora-finetune_2bs_1node_1gpu_1000s \ --lr=1.5e-05 \ --min-lr=1.49e-05 \ --warmup-steps=10 \ --result-dir=/jet/logs/recipe/model-evo2_variant-lora-finetune_bionemo_partial-conv_nodes-1_gpus-1_batch_size-2_config_name-1b_max_steps-1000_stop_steps-100_target-dgxh100-eos_task-lora-finetune-from-ckpt/tensorboard_logs \ --micro-batch-size=2 \ --grad-acc-batches=4 \ --limit-val-batches=20 \ --seq-length=8192 \ --clip-grad=250 \ --wd=0.001 \ --attention-dropout=0.01 \ --hidden-dropout=0.01 \ --num-layers 4 \ --hybrid-override-pattern 'SDH*' \ --devices=1 \ --num-nodes=1 \ --val-check-interval=5 \ --wandb-project=jet--partial-conv--main \ --wandb-group=evo2_lora_finetune_1b_lora_finetune_from_ckpt_dgxh100_eos \ --create-tensorboard-logger --activation-checkpoint-recompute-num-layers=2 \ --disable-checkpointing \ --early-stop-on-step=100 \ --garbage-collect-at-inference

Error Messages and Logs

Training logs 

......
0: Training epoch 0, iteration 0/99 | lr: 0 | global_batch_size: 8 | global_step: 0 | reduced_train_loss: 6.278 | train_step_timing in s: 3.453
0: Training epoch 0, iteration 1/99 | lr: 1.5e-06 | global_batch_size: 8 | global_step: 1 | reduced_train_loss: 6.3 | train_step_timing in s: 0.1276 | consumed_samples: 16
0: Training epoch 0, iteration 2/99 | lr: 3e-06 | global_batch_size: 8 | global_step: 2 | reduced_train_loss: 6.283 | train_step_timing in s: 0.1269 | consumed_samples: 24
0: Training epoch 0, iteration 3/99 | lr: 4.5e-06 | global_batch_size: 8 | global_step: 3 | reduced_train_loss: 6.285 | train_step_timing in s: 0.1271 | consumed_samples: 32
0: Training epoch 0, iteration 4/99 | lr: 6e-06 | global_batch_size: 8 | global_step: 4 | reduced_train_loss: 6.296 | train_step_timing in s: 0.1269 | consumed_samples: 40
0: [NeMo W 2025-09-06 12:29:56 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:431: It is recommended to use `self.log('global_batch_size', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.
0:     
0: [NeMo W 2025-09-06 12:29:56 nemo_logging:405] /usr/local/lib/python3.12/dist-packages/lightning/pytorch/trainer/connectors/logger_connector/result.py:431: It is recommended to use `self.log('val_loss', ..., sync_dist=True)` when logging on epoch level in distributed setting to accumulate the metric across devices.
0:     
0: Validation: iteration 1/20
0: Validation: iteration 2/20
0: Validation: iteration 3/20
0: Validation: iteration 4/20
0: Validation: iteration 5/20
0: Validation: iteration 6/20
0: Validation: iteration 7/20
0: Validation: iteration 8/20
0: Validation: iteration 9/20
0: Validation: iteration 10/20
0: Validation: iteration 11/20
0: Validation: iteration 12/20
0: Validation: iteration 13/20
0: Validation: iteration 14/20
0: Validation: iteration 15/20
0: Validation: iteration 16/20
0: Validation: iteration 17/20
0: Validation: iteration 18/20
0: Validation: iteration 19/20
0: Validation: iteration 20/20
0: Training epoch 0, iteration 5/99 | lr: 7.5e-06 | global_batch_size: 8 | global_step: 5 | reduced_train_loss: 6.299 | train_step_timing in s: 0.1278 | consumed_samples: 48 | val_loss: 6.285
0: Training epoch 0, iteration 6/99 | lr: 9e-06 | global_batch_size: 8 | global_step: 6 | reduced_train_loss: 6.282 | train_step_timing in s: 0.127 | consumed_samples: 56 | val_loss: 6.285
0: Training epoch 0, iteration 7/99 | lr: 1.05e-05 | global_batch_size: 8 | global_step: 7 | reduced_train_loss: 6.293 | train_step_timing in s: 0.127 | consumed_samples: 64 | val_loss: 6.285
0: Training epoch 0, iteration 8/99 | lr: 1.2e-05 | global_batch_size: 8 | global_step: 8 | reduced_train_loss: 6.288 | train_step_timing in s: 0.1271 | consumed_samples: 72 | val_loss: 6.285
0: Training epoch 0, iteration 9/99 | lr: 1.35e-05 | global_batch_size: 8 | global_step: 9 | reduced_train_loss: 6.281 | train_step_timing in s: 0.1272 | consumed_samples: 80 | val_loss: 6.285
0: Validation: iteration 1/20
0: Validation: iteration 2/20
0: Validation: iteration 3/20
0: Validation: iteration 4/20
0: Validation: iteration 5/20
0: Validation: iteration 6/20
0: Validation: iteration 7/20
0: Validation: iteration 8/20
0: Validation: iteration 9/20
0: Validation: iteration 10/20
0: Validation: iteration 11/20
0: Validation: iteration 12/20
0: Validation: iteration 13/20
0: Validation: iteration 14/20
0: Validation: iteration 15/20
0: Validation: iteration 16/20
0: Validation: iteration 17/20
0: Validation: iteration 18/20
0: Validation: iteration 19/20
0: Validation: iteration 20/20
0: Training epoch 0, iteration 10/99 | lr: 1.5e-05 | global_batch_size: 8 | global_step: 10 | reduced_train_loss: 6.296 | train_step_timing in s: 0.1277 | consumed_samples: 88 | val_loss: 6.285
0: Training epoch 0, iteration 11/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 11 | reduced_train_loss: 6.283 | train_step_timing in s: 0.127 | consumed_samples: 96 | val_loss: 6.285
0: Training epoch 0, iteration 12/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 12 | reduced_train_loss: 6.265 | train_step_timing in s: 0.127 | consumed_samples: 104 | val_loss: 6.285
0: Training epoch 0, iteration 13/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 13 | reduced_train_loss: 6.275 | train_step_timing in s: 0.1268 | consumed_samples: 112 | val_loss: 6.285
0: Training epoch 0, iteration 14/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 14 | reduced_train_loss: 6.295 | train_step_timing in s: 0.1269 | consumed_samples: 120 | val_loss: 6.285
0: Validation: iteration 1/20
0: Validation: iteration 2/20
0: Validation: iteration 3/20
0: Validation: iteration 4/20
0: Validation: iteration 5/20
0: Validation: iteration 6/20
0: Validation: iteration 7/20
0: Validation: iteration 8/20
0: Validation: iteration 9/20
0: Validation: iteration 10/20
0: Validation: iteration 11/20
0: Validation: iteration 12/20
0: Validation: iteration 13/20
0: Validation: iteration 14/20
0: Validation: iteration 15/20
0: Validation: iteration 16/20
0: Validation: iteration 17/20
0: Validation: iteration 18/20
0: Validation: iteration 19/20
0: Validation: iteration 20/20
0: Training epoch 0, iteration 15/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 15 | reduced_train_loss: 6.287 | train_step_timing in s: 0.1276 | consumed_samples: 128 | val_loss: 6.285
0: Training epoch 0, iteration 16/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 16 | reduced_train_loss: 6.284 | train_step_timing in s: 0.1271 | consumed_samples: 136 | val_loss: 6.285
0: Training epoch 0, iteration 17/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 17 | reduced_train_loss: 6.31 | train_step_timing in s: 0.1267 | consumed_samples: 144 | val_loss: 6.285
0: Training epoch 0, iteration 18/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 18 | reduced_train_loss: 6.28 | train_step_timing in s: 0.1268 | consumed_samples: 152 | val_loss: 6.285
0: Training epoch 0, iteration 19/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 19 | reduced_train_loss: 6.29 | train_step_timing in s: 0.1268 | consumed_samples: 160 | val_loss: 6.285
0: Validation: iteration 1/20
0: Validation: iteration 2/20
0: Validation: iteration 3/20
0: Validation: iteration 4/20
0: Validation: iteration 5/20
0: Validation: iteration 6/20
0: Validation: iteration 7/20
0: Validation: iteration 8/20
0: Validation: iteration 9/20
0: Validation: iteration 10/20
0: Validation: iteration 11/20
0: Validation: iteration 12/20
0: Validation: iteration 13/20
0: Validation: iteration 14/20
0: Validation: iteration 15/20
0: Validation: iteration 16/20
0: Validation: iteration 17/20
0: Validation: iteration 18/20
0: Validation: iteration 19/20
0: Validation: iteration 20/20
0: Training epoch 0, iteration 20/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 20 | reduced_train_loss: 6.274 | train_step_timing in s: 0.1278 | consumed_samples: 168 | val_loss: 6.285
0: Training epoch 0, iteration 21/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 21 | reduced_train_loss: 6.298 | train_step_timing in s: 0.1271 | consumed_samples: 176 | val_loss: 6.285
0: Training epoch 0, iteration 22/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 22 | reduced_train_loss: 6.284 | train_step_timing in s: 0.1269 | consumed_samples: 184 | val_loss: 6.285
0: Training epoch 0, iteration 23/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 23 | reduced_train_loss: 6.281 | train_step_timing in s: 0.1269 | consumed_samples: 192 | val_loss: 6.285
0: Training epoch 0, iteration 24/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 24 | reduced_train_loss: 6.292 | train_step_timing in s: 0.127 | consumed_samples: 200 | val_loss: 6.285
0: Validation: iteration 1/20
0: Validation: iteration 2/20
0: Validation: iteration 3/20
0: Validation: iteration 4/20
0: Validation: iteration 5/20
0: Validation: iteration 6/20
0: Validation: iteration 7/20
0: Validation: iteration 8/20
0: Validation: iteration 9/20
0: Validation: iteration 10/20
0: Validation: iteration 11/20
0: Validation: iteration 12/20
0: Validation: iteration 13/20
0: Validation: iteration 14/20
0: Validation: iteration 15/20
0: Validation: iteration 16/20
0: Validation: iteration 17/20
0: Validation: iteration 18/20
0: Validation: iteration 19/20
0: Validation: iteration 20/20
0: Training epoch 0, iteration 25/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 25 | reduced_train_loss: 6.293 | train_step_timing in s: 0.1274 | consumed_samples: 208 | val_loss: 6.285
0: Training epoch 0, iteration 26/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 26 | reduced_train_loss: 6.274 | train_step_timing in s: 0.1269 | consumed_samples: 216 | val_loss: 6.285
0: Training epoch 0, iteration 27/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 27 | reduced_train_loss: 6.295 | train_step_timing in s: 0.1269 | consumed_samples: 224 | val_loss: 6.285
0: Training epoch 0, iteration 28/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 28 | reduced_train_loss: 6.29 | train_step_timing in s: 0.1268 | consumed_samples: 232 | val_loss: 6.285
0: Training epoch 0, iteration 29/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 29 | reduced_train_loss: 6.26 | train_step_timing in s: 0.1266 | consumed_samples: 240 | val_loss: 6.285
0: Validation: iteration 1/20
0: Validation: iteration 2/20
0: Validation: iteration 3/20
0: Validation: iteration 4/20
0: Validation: iteration 5/20
0: Validation: iteration 6/20
0: Validation: iteration 7/20
0: Validation: iteration 8/20
0: Validation: iteration 9/20
0: Validation: iteration 10/20
0: Validation: iteration 11/20
0: Validation: iteration 12/20
0: Validation: iteration 13/20
0: Validation: iteration 14/20
0: Validation: iteration 15/20
0: Validation: iteration 16/20
0: Validation: iteration 17/20
0: Validation: iteration 18/20
0: Validation: iteration 19/20
0: Validation: iteration 20/20
0: Training epoch 0, iteration 30/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 30 | reduced_train_loss: 6.277 | train_step_timing in s: 0.1276 | consumed_samples: 248 | val_loss: 6.285
0: Training epoch 0, iteration 31/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 31 | reduced_train_loss: 6.27 | train_step_timing in s: 0.1271 | consumed_samples: 256 | val_loss: 6.285
0: Training epoch 0, iteration 32/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 32 | reduced_train_loss: 6.288 | train_step_timing in s: 0.1271 | consumed_samples: 264 | val_loss: 6.285
0: Training epoch 0, iteration 33/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 33 | reduced_train_loss: 6.28 | train_step_timing in s: 0.1268 | consumed_samples: 272 | val_loss: 6.285
0: Training epoch 0, iteration 34/99 | lr: 1.49e-05 | global_batch_size: 8 | global_step: 34 | reduced_train_loss: 6.3 | train_step_timing in s: 0.1268 | consumed_samples: 280 | val_loss: 6.285

Docker Image

No response

System Information

Environment Details:

• OS: ubuntu
• CPU: amd64

GPU Details:

• GPU Model: H100

Additional Context

No response

[cclough]

Beyond the loss decreasing, we are also missing functional proof that fine-tuning works, with or without LoRa. It would be helpful to add an inference demo to the fine-tuning notebook, showing e.g. an improvement in sequence completion % when generating the second half of a gene which has been fine-tuned on, vs base model.

cc @jstjohn @spoulos-nvidia

[balvisio]

@cclough : We have made progress making LoRA finetuning work with Evo2. We will have a fixed version in the coming weeks.

[balvisio]

@cclough : This PR #1335 contains a working version of Evo2LoRA.
Note:

1. I have validated LoRA fine-tuning a 1B Evo2 achieves same val loss as full fine-tuning. However, this branch is still experimental.
2. This branch won't work with activation checkpointing. Make sure to pass the flag --no-activation-checkpointing in the Evo2 training command. I will work on a fix in the coming weeks.

[jannis-baum]

@balvisio any updates on this? I tried with the base version 2.7.1 and with your #1335 but either way it doesn't seem to be working, the loss stays quite constant.

I also saw that you've made some very big refactors since the last release, are there any plans on having a stable version out soon?

[balvisio]

@jannis-baum We highly recommend you to use the Megratron-bridge based Evo2 recipe instead of the previous implementation in sub-packages/. We recently added LoRA support: docs and LoRA Jupyter Notebook demo. Let me know if you run into any issue.

[jannis-baum]

@balvisio Thanks for the quick reply! I'll try that. Should I use the nvcr.io/nvidia/clara/bionemo-framework:nightly container then or is there some fixed version you would recommend? The page linked in the README (https://docs.nvidia.com/bionemo-framework/latest/user-guide/) doesn't seem to exist anymore

[balvisio]

@jannis-baum Each Bionemo recipes can be built independently. For Evo2 recipe see here.

[jannis-baum]

@jannis-baum Each Bionemo recipes can be built independently. For Evo2 recipe see here.

Thank you, that works! One note: When using train_evo2 --lora-finetune and tensorboard logging I get a crash because report_l2_norm_grad in megatron-bridge's train_utils.py iterates over all model params and accesses p.main_grad unconditionally but with LoRA the base model is frozen, i.e. there's no main_grad allocated, and an error is raised.

I was able to resolve it by monkey-patching the report_l2_norm_grad function with this wrapper around train_evo2, basically just adding p.main_grad = None so it's skipped (as supported by the original function) rather than crashing:

import sys

import megatron.bridge.training.utils.train_utils as _tu

_orig_report_l2_norm_grad = _tu.report_l2_norm_grad


def _report_l2_norm_grad_lora_safe(model):
    modules = model if isinstance(model, list) else [model]
    for m in modules:
        for p in m.parameters():
            if not hasattr(p, "main_grad"):
                p.main_grad = None
    return _orig_report_l2_norm_grad(model)


_tu.report_l2_norm_grad = _report_l2_norm_grad_lora_safe

if __name__ == "__main__":
    from bionemo.evo2.run.train import main  # noqa: E402
    sys.exit(main())

[balvisio]

@jannis-baum Thank you. Yes, there is a fix in Megatron-Bridge for that issue here. Once we validate the new release of Megatron-Bridge we will upgrade the library in the recipe to get that issue fixed.