Add a classification / regression finetuning head to evo2

[MidAtBest]

Description

I am struggling to understand if there is an easy to finetune a Hyena model (EVO2) on a sequence leve ( e.g. ATCGATCG -> 1) and token level classification task ( e.g. ATCGATCG -> 10001000)

If so I am not sure what are the modules I should add / modify to allow for this.

Solution desired

Having a sequence / token level finetuning strategy to choose from or a solution to implement it on my side.

Any help on this issue would be very much appreciated.

[dorotat-nv]

@MidAtBest , please provide commit sha or version tag for your request

[MidAtBest]

@dorotat-nv sorry for the delay in my response, here is the commit sha 495ba03

[jstjohn]

@MidAtBest you can subclass the hyena model and modify the loss function in the forward pass similar to what we did in the mamba model here? https://github.com/NVIDIA/bionemo-framework/blob/main/sub-packages/bionemo-evo2/src/bionemo/evo2/models/mamba.py#L250-L263. The megatron level loss will try to get the num_tokens from your loss mask. So one thing you can do is just return torch.ones(batch_size, 1) for the loss mask in the case of your sequence -> 1 output fine-tuning task.

When you restore model weights from a checkpoint, one issue is that new heads will show up as new layer weights, so you'll need to write a checkpoint converter that maps old weights to new weights. Basically what you'll do is load old weights where possible (all layers other than your new head) and then use the newly initialized weights for the new head by not pulling anything in for that. https://github.com/NVIDIA/bionemo-framework/blob/main/sub-packages/bionemo-evo2/src/bionemo/evo2/utils/checkpoint/convert_to_nemo.py#L63-L158 can be a starting point for that. This just removes the optimizer from a saved checkpoint.

So high level your flow would look like this:

0. Create a new subpackage by copying the bionemo-evo2 subpackage and delete everything you don't need. Basically clean it out for your project but use that as a template.
1. Subclass the hyena model (and the hyena gpt model) like we do for the mamba model in the example above. That lives in nemo/collections/llm/gpt/hyena.py but you can just import the parent from there and create your new configs/subclasses in your package.
2. In your subclasses do one for the different loss types you want, or if you want multiple losses that's fine too, just add all of them to forward along with your preferred way of collapsing, determining which to apply based on the contents of the batch, etc.
3. If you do not need model parallelism, 2 is easy. If you need model parallelism you'll need to think about proper reductions of your loss across the tensor parallel groups for example. That can be a bit tricky to get right. You'll need to modify your forward/backwards passes potentially to do proper reductions across the tensor parallel group. See for example https://huggingface.co/spaces/nanotron/ultrascale-playbook on some background if you don't know what i'm talking about. https://github.com/NVIDIA/bionemo-framework/blob/main/sub-packages/bionemo-evo2/src/bionemo/evo2/utils/loss/embedding_variance.py#L27-L192 is an example of a loss where we did that. The better ai coding agents can probably help reasonably with that task, but definitely make some tests to verify that the backwards works as intended in a parallel setup.
4. If you write custom dataloaders note the requirements that megatron has for dataloaders in a parallel environment. Basically they need to be deterministic. You need the property that getitem(i) from your dataloader always returns the same thing. You also need your dataloader to handle it's own shuffling internally (and again deterministically). So like make use of seeds etc inside of your loaders. If you need to handle concepts like epochs, we have some examples in bionemo that take care of mapping getitem(i) -> epoch_j, element_k so you can make a seed based on the two integers j,k if necessary, so you do not have the same seed applied to your data every time you resample element k, if that makes sense.

Feel free to ask follow up questions if anything is not clear or you hit roadblocks trying this out.

[MidAtBest]

Hi @jstjohn thanks a lot for the answer will give it a go and update here if any questions arise (or if I succeed in carrying this out as it might be helpful to other users).