Evo2 pompt generation

[michele1993]

Hello,

I have been trying to pass different prompts to the infer_evo2 script using the --prompt flag, however no matter what prompts I pass to the model, the model seems to generate nearly identical sequences, as if the passed prompts had no effect on the model generation. I have tried this for both the 1-billion 8k and 7-billion 8k base models as well as analogous versions fine-tuned on a custom dataset, testing different prompts and sequence lengths.

Example
If I pass no prompt, which defaults to the default prompt (i.e., "|d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;..." etc.), by running

infer_evo2 --ckpt-dir $BIONEMO_CACHE/nemo2_evo2_7b_8k/ --temperature 0.7 --top-k 4 --max-new-tokens 100

I get the following output sequence,
"
CTCTTCTGGTATTTGGTGTGCGATGGACCTATTGAGCTGATAGTGTGGATCCTGTTGCTCCGGATGGTTTTGCTGGCCAACGGCCCCGCAACGCGTGCAC
"
If I prompt the model with a very specific prompt of an unrelated gene to the default prompt,

infer_evo2 --ckpt-dir $BIONEMO_CACHE/nemo2_evo2_7b_8k/ --temperature 0.7 --top-k 4 --max-new-tokens 1000 --prompt GAGCCTCCTTAACTCTTTACCGTTCACGTTAAATCACTTTCGCTACCGCGTAAACGTGCACGAGCCACTATAGTGGCGC

The model outputs a sequences that is over 90% identical to the sequence outputted with the default prompt, suggesting the provided prompt had no effect on the outputted sequence,
"
GTCTTGCGGTATTTTGTGTGCGATGGACCTATTGAGCGAATAGTGTGGATCCTGTTGCTCCGGATGGTTT
GACTGGCCAACGGCCCCCCAACGCGGGCAC
"

I have tried this for different prompts, different temperature or top-k values, always observing the same behaviour.
Additionally, during the inference process, the following warning is outputted six times,
"
[WARNING | py.warnings ]: /usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/attention.py:5213: UserWarning: window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=no_mask warnings.warn(
"
I am not sure if I am prompting the model incorrectly or this reflects some deeper issue. Any help with this would be greatly appreciated, thanks.

[trvachov]

@michele1993 thanks for reporting the issue, sounds like a bug. @jwilber can you triage?

[jwilber]

Thanks for the issue. 90% identical seems to infer that prompting is having an effect, though maybe we should expect more.

@tshimko-nv Can you check this out? I wonder if it's related to the particular values of the prompts being used here?

[tshimko-nv]

Hi @michele1993 -

Thank you for reporting this issue with the infer_evo2 script. I can confirm that we are seeing the same behavior you've described.

I am actively looking into this issue and have been able to reproduce your results. Using my own testing script, I've confirmed that different prompts are indeed producing sequences with unexpectedly high similarity scores.

My Test Results

I ran two test scenarios:

1. Test Set 1: Mixed Different Prompts
   • Used various different prompts (taxonomic information + DNA sequences)
   • Generated sequences showed identity scores ranging from ~89% to 99%
   • This high similarity suggests the prompts are not having the expected differentiation effect
2. Test Set 2: Identical Prompts (Control)
   • Used identical prompts for comparison
   • As expected, generated sequences showed 100% identity

In short - the prompts are having an effect, but I agree that it's not as large an affect as we might naively expect.

==========================================
IDENTITY MATRIX FOR PROMPT SET 1 (Mixed prompts)
==========================================
Matrix shows pairwise sequence identity (lower triangle):
Each row/column represents one generated sequence

 1: 1.000 
 2: 0.900 1.000 
 3: 0.970 0.910 1.000 
 4: 0.970 0.910 0.940 1.000 
 5: 0.950 0.900 0.940 0.930 1.000 
 6: 0.900 0.980 0.930 0.890 0.890 1.000 
 7: 0.930 0.940 0.960 0.910 0.910 0.950 1.000 
 8: 0.960 0.900 0.950 0.930 0.990 0.900 0.910 1.000 
 9: 0.970 0.910 0.960 0.940 0.980 0.910 0.920 0.990 1.000 
10: 0.980 0.900 0.970 0.950 0.970 0.900 0.930 0.980 0.990 1.000 
11: 0.930 0.940 0.960 0.910 0.910 0.950 1.000 0.910 0.920 0.930 1.000 


==========================================
IDENTITY MATRIX FOR PROMPT SET 2 (Identical prompts)
==========================================
Matrix shows pairwise sequence identity (lower triangle):
Each row/column represents one generated sequence

 1: 1.000 
 2: 1.000 1.000 
 3: 1.000 1.000 1.000 
 4: 1.000 1.000 1.000 1.000 
 5: 1.000 1.000 1.000 1.000 1.000 
 6: 1.000 1.000 1.000 1.000 1.000 1.000 
 7: 1.000 1.000 1.000 1.000 1.000 1.000 1.000 
 8: 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 
 9: 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 
10: 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 1.000 

==========================================
COMPARISON SUMMARY
==========================================
Set 1: Mixed different prompts (taxonomic info + DNA sequences)
Set 2: Identical prompts (same DNA sequence repeated)

Expected results:
- Set 1 should show varied identity scores due to different inputs
- Set 2 should show high identity scores due to identical inputs

The results from Test Set 1 align with your observations - different prompts should produce more varied sequences, but instead we're seeing high identity scores (90%+) similar to what you reported. I'll update once I have further information. In the meantime, have you adjusting any of the sampling params as well (temperature, top-k, etc.)?

[michele1993]

Hi all,

Thank you for getting back to me.

I have just run some experiments comparing two sequences generated for different top-p ([0.1, 0.6, 1]) and temperature ([0.2, 0.5, 1]) values. I have done this using the same prompt for both sequences or using two (very) different prompts.

The two prompts I used are:

--prompt "|d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Escherichia;s__Escherichia|"

and

--prompt "|D__ANIMALIA;P__CHORDATA;C__MAMMALIA;O__DIPROTODONTIA;F__PHASCOLARCTIDAE;G__PHASCOLARCTOS;S__PHASCOLARCTOS CINEREUS|"

which should refer to different species. When I generated both sequences with the same prompt, I just used the first one (i.e., default prompt).

Below I included a table with the sequence identity for pairs of sequences generated across the different top-p and temperature values for both same prompt and different prompt generation,

====================
Pair of sequences generated with the same  top-p (across different temperature values)
=====================
Both sequences generated at  top-p: 0.1
Temperature values of the two sequences (0.2,0.2) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Temperature values of the two sequences (0.2,0.5) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Temperature values of the two sequences (0.2,1) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Temperature values of the two sequences (0.5,0.5) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Temperature values of the two sequences (0.5,1) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Temperature values of the two sequences (1,1) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)

Both sequences generated top-p: 0.6
Temperature values of the two sequences (0.2,0.2) | Similarity: 100.0 (same prompt) & 88.57421875 (different prompt)
Temperature values of the two sequences (0.2,0.5) | Similarity: 77.24609375 (same prompt) & 77.24609375 (different prompt)
Temperature values of the two sequences (0.2,1) | Similarity: 71.2890625 (same prompt) & 70.5078125 (different prompt)
Temperature values of the two sequences (0.5,0.5) | Similarity: 100.0 (same prompt) & 88.18359375 (different prompt)
Temperature values of the two sequences (0.5,1) | Similarity: 90.4296875 (same prompt) & 87.20703125 (different prompt)
Temperature values of the two sequences (1,1) | Similarity: 100.0 (same prompt) & 92.3828125 (different prompt)

Both sequences generated Top-p: 1
Temperature values of the two sequences (0.2,0.2) | Similarity: 100.0 (same prompt) & 95.41015625 (different prompt)
Temperature values of the two sequences (0.2,0.5) | Similarity: 89.453125 (same prompt) & 89.0625 (different prompt)
Temperature values of the two sequences (0.2,1) | Similarity: 86.5234375 (same prompt) & 86.5234375 (different prompt)
Temperature values of the two sequences (0.5,0.5) | Similarity: 100.0 (same prompt) & 98.14453125 (different prompt)
Temperature values of the two sequences (0.5,1) | Similarity: 96.484375 (same prompt) & 96.58203125 (different prompt)
Temperature values of the two sequences (1,1) | Similarity: 100.0 (same prompt) & 98.92578125 (different prompt)

=====================
Pair of sequences generated with the same temperature  (across different top-p values)
=====================
Both sequences generated at temperature: 0.2
Top-p value of the two sequences (0.1,0.1) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Top-p value of the two sequences (0.1,0.6) | Similarity: 36.03515625 (same prompt) & 34.86328125 (different prompt)
Top-p value of the two sequences (0.1,1) | Similarity: 28.90625 (same prompt) & 30.2734375 (different prompt)
Top-p value of the two sequences (0.6,0.6) | Similarity: 100.0 (same prompt) & 88.57421875 (different prompt)
Top-p value of the two sequences (0.6,1) | Similarity: 71.875 (same prompt) & 71.09375 (different prompt)
Top-p value of the two sequences (1,1) | Similarity: 100.0 (same prompt) & 95.41015625 (different prompt)

Both sequences generated at temperature: 0.5
Top-p value of the two sequences (0.1,0.1) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Top-p value of the two sequences (0.1,0.6) | Similarity: 26.26953125 (same prompt) & 29.78515625 (different prompt)
Top-p value of the two sequences (0.1,1) | Similarity: 26.66015625 (same prompt) & 27.1484375 (different prompt)
Top-p value of the two sequences (0.6,0.6) | Similarity: 100.0 (same prompt) & 88.18359375 (different prompt)
Top-p value of the two sequences (0.6,1) | Similarity: 76.46484375 (same prompt) & 75.87890625 (different prompt)
Top-p value of the two sequences (1,1) | Similarity: 100.0 (same prompt) & 98.14453125 (different prompt)

Both sequences generated at temperature: 1
Top-p value of the two sequences (0.1,0.1) | Similarity: 100.0 (same prompt) & 99.90234375 (different prompt)
Top-p value of the two sequences (0.1,0.6) | Similarity: 24.12109375 (same prompt) & 26.26953125 (different prompt)
Top-p value of the two sequences (0.1,1) | Similarity: 25.9765625 (same prompt) & 26.66015625 (different prompt)
Top-p value of the two sequences (0.6,0.6) | Similarity: 100.0 (same prompt) & 92.3828125 (different prompt)
Top-p value of the two sequences (0.6,1) | Similarity: 79.1015625 (same prompt) & 78.61328125 (different prompt)
Top-p value of the two sequences (1,1) | Similarity: 100.0 (same prompt) & 98.92578125 (different prompt)

I am not too sure what to make of these. I guess you get the greatest difference in sequence identity when the two sequences were generated by the most different parameters (e.g., one sequence generated with top-p=1 and the other with top-p=0.1 etc.), as you would expect. This is also true when two different prompts were used, in few cases, even leading to two sequences with slightly higher sequence identity when two different prompts were used compared to using the same prompt.

[tshimko-nv]

Thanks for this thorough analysis! I'll keep digging on this bug today and will update here with what I can determine.

[tshimko-nv]

Hi @michele1993 - Just wanted to provide a quick update here: the team has identified the root cause of this bug (incompatibility between the Megatron-based implementation in BioNeMo and the original Vortex kernels used for generation) and we're actively working on a fix. Unfortunately, the fix will rely on updates to Megatron + NeMo. We'll try to update here as progress is made, but we're unable to provide a timeline for the fix at this time.

[michele1993]

Hi @tshimko-nv thank you for the update. I understand, happy to know you could find the root cause. Is there any way to convert a (fine-tuned) nemo evo2 model back to Savanna and use their generation API ?

[tshimko-nv]

Hi @michele1993 - unfortunately we don't currently have a way to convert between the NeMo checkpoint and the Vortex inference API. We would also note that Savanna is the training interface for Evo2, Vortex is the framework that can handle inference. So ideally you would want to convert between the BioNeMo/NeMo checkpoint to Vortex, not to Savanna.

I think your best bet is to wait on the in-progress fix to make it into a branch here, then you can build the container and/or mount the working code into the nightly container build to run inference locally instead of needing to wait for the next release. It should be easier for us to get the tests related to inference passing on a branch rather than into main/released container. If you're willing to build the container/mount the code yourself and deal with the extra steps involved there, that's probably the quickest solution. @jwilber is the engineer currently working on that fix, so I'll tag him here to provide updates as that work progresses.

[jwilber]

Hey @michele1993 - thanks for your patientence here. It will take some time to get these changes on main, but I've created a branch you can use in the meantime: #949
The branch has the code that fixes the issue we described previously: incompatibility between the Megatron-based implementation in BioNeMo and the original Vortex kernels used for generation. Given this discussion, the branch also has a new test that shows the expected effect of different prompts on generate functionality.

To apply these changes easily, check out the jwilber/evo2-generate-patch branch and update the submodules:

git switch jwilber/evo2-generate-patch
git submodule update --init --recursive

If you're using the nightly container, you can simply mount the updated repositories to access the changes. For example, when I tested this on a Slurm cluster using the nightly container, I first ran the commands above, then mounted the local repositories.
(Note that you'll also have to update your PYTHONPATH to use these as well, see below).

For example, here's a basic outline of the .srun script that I used to test these changes myself:

#!/bin/bash
IMG=nvcr.io/nvidia/clara/bionemo-framework:nightly

DATA_PATH=.../data/loc/data
DATA_MOUNT=/workspace/bionemo2/data
EVO2_MOUNT=/workspace/bionemo2/sub-packages/bionemo-evo2
EVO2_PATH=.../jwilber/evo2/bionemo-framework/sub-packages/bionemo-evo2
THIRDPARTY_PATH=.../jwilber/evo2/bionemo-framework/3rdparty
THIRDPARTY_MOUNT=/workspace/bionemo2/3rdparty

EVO2_SRC_MOUNT=${EVO2_MOUNT}/src

MOUNTS="${DATA_PATH}:${DATA_MOUNT},\
${EVO2_PATH}:${EVO2_MOUNT},\
${THIRDPARTY_PATH}:${THIRDPARTY_MOUNT}"

export PYTHONPATH="${CORE_SRC_MOUNT}:${EVO2_SRC_MOUNT}:${THIRDPARTY_MOUNT}/Megatron-LM:${THIRDPARTY_MOUNT}/NeMo:${PYTHONPATH:-}"

srun -A account \
     -N1 \
     -p interactive \
     -J account-... \
     --container-image="$IMG" \
     --container-mounts="$MOUNTS" \
     --container-env=PYTHONPATH="$PYTHONPATH" \
     --pty bash

Hope that helps! Do let me know if you find any continued oddities here. Also, if you don't have fp8 available on your accelerators, you'll want to toggle this to False

[michele1993]

Hi @jwilber and @tshimko-nv

Thank you so much for the patch and the super helpful guidance.

UPDATE:
I think I am still observing some strange behaviors.

First, I tested the model with the exact same prompts and settings in test script and found sequences with 0.41 similarity across the two different prompts (not sure this replicas what you observed).

Second, I increased the number of maximum tokens to be generated (to 500) and I obtained more similar sequences across the two different default prompts compared to the shorter generation (0.586 similarity with BLAST finding 88% match with 27%). This is strange as I would expect longer sequences to be more different if we are prompting from different prompts.

Third I tested the model by generating "long" sequences (2500 tokens) across 3 different prompts:
1)

|d__Bacteria;p__Pseudomonadota;c__Gammaproteobacteria;o__Enterobacterales;f__Enterobacteriaceae;g__Escherichia;s__Escherichia|

 |D__ANIMALIA;P__CHORDATA;C__MAMMALIA;O__DIPROTODONTIA;F__PHASCOLARCTIDAE;G__PHASCOLARCTOS;S__PHASCOLARCTOS CINEREUS|

GAGCCTCCTTAACTCTTTACCGTTCACGTTAAATCACTTTCGCTACCGCGTAAACGTGCACGAGCCACTATAGTGGCGCGAAAAAGACTCCATTTCAACT

All these prompts produced highly similar ~85% in BLAST sequences. I find this fairly unlikely as these are different species (not an expert though). Interestingly, when I reduced the MAX_NEW_TOKENS I got more dissimilar sequences.

I will try to run more tests, by looking for two species with DNA that is for sure very different.

Quick question:
What should CORE_SRC_MOUNT point to ? should it point to bionemo-framework/sub-packages/bionemo-core/src from jwilber/evo2-generate-patch ? and if so, shouldn't this also be mounted ? (I think I am getting away without defining it properly, but I am scared it may lead to unexpected behaviors)

Thank you so much again :)