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).

[MidAtBest]

Thank you for your help, this has been very useful!

Currently I am struggling with saving to S3, the best solution I have found it to mirror the checkpoints directly to S3 using the callback below.

code

class MirrorDistCkptToS3(Callback):
    """
    A PyTorch Lightning callback that mirrors finalized distributed checkpoints
    (produced by NeMo/Megatron) from the local filesystem to an S3 bucket using fsspec.

    ### Purpose
    NeMo's distributed checkpointing system (`dist_checkpointing`) saves each
    checkpoint as a *directory* (e.g., `epoch=0-step=100-last/weights/`) locally.
    However, when training on clusters or cloud VMs, these directories often
    reside on ephemeral storage and are not automatically synced to S3.

    This callback watches for new local distributed checkpoints produced by
    NeMo’s `ModelCheckpoint` callback, uploads the fully written checkpoint
    directory to S3, and prunes any older remote checkpoints that are no longer
    present locally (keeping S3 in sync with your local `save_top_k` policy).

    ### Key Features
    - Automatically uploads new distributed checkpoints to S3 using fsspec.
    - Only mirrors finalized checkpoints (ignores partially written ones).
    - Deletes any remote checkpoints that were pruned locally (top-k retention).
    - Runs only on the global rank 0 process (to avoid duplicate uploads).
    - Works transparently with NeMo’s `ckpt_to_dir()` mapping for distributed ckpts.

    ### Usage Example
    ```python
    checkpoint_cb = nl_callbacks.ModelCheckpoint(
        dirpath="./results/checkpoints",
        save_last=True,
        save_top_k=3,
        monitor="val_loss",
    )

    mirror_cb = MirrorDistCkptToS3(
        s3_prefix="s3://mybucket/my-exp/checkpoints",
        src_root="./results/checkpoints",
        checkpoint_cb=checkpoint_cb,
    )

    trainer = nl.Trainer(callbacks=[checkpoint_cb, mirror_cb, ...])
    ```

    After each validation or training step, this callback ensures your
    `s3://mybucket/my-exp/checkpoints` folder reflects your local top-k checkpoints.

    Args:
        s3_prefix (str): The S3 prefix to mirror checkpoints under
            (e.g., "s3://mybucket/path/to/checkpoints").
        src_root (str): The local directory where NeMo writes distributed checkpoints.
        checkpoint_cb (nl_callbacks.ModelCheckpoint): The active NeMo checkpoint callback.
    """

    def __init__(self, s3_prefix: str, src_root: str, checkpoint_cb: nl_callbacks.ModelCheckpoint):
        self.s3_prefix = s3_prefix.rstrip("/")
        self.src_root = str(src_root).rstrip("/")
        self.cb = checkpoint_cb
        self._last_mirrored_local_dir = None
        self.fs = fsspec.filesystem("s3")  # s3fs-based fsspec filesystem client

    # ---------------------------------------------------------------------
    # Utility helpers
    # ---------------------------------------------------------------------

    def _is_global_zero(self, trainer) -> bool:
        """Only rank 0 should perform S3 sync to avoid redundant uploads."""
        return getattr(trainer, "is_global_zero", True)

    def _is_finalized(self, local_dir: str) -> bool:
        """
        Check whether a checkpoint directory is finalized.

        NeMo creates an "unfinished" marker file while a checkpoint is being
        written. This method checks for its absence to confirm finalization.
        """
        marker = f"{local_dir}{UNFINISHED_SUFFIX}"
        return not os.path.exists(marker)

    def _upload_dir(self, local_dir: str, remote_dir: str):
        """
        Recursively upload a directory to S3.

        Walks the local directory tree, computing relative paths and using
        fsspec’s `.put()` to copy each file to the remote location.
        """
        for root, _, files in os.walk(local_dir):
            rel = os.path.relpath(root, local_dir)
            rel = "" if rel == "." else rel
            dst_dir = f"{remote_dir}/{rel}" if rel else remote_dir
            for fn in files:
                lp = os.path.join(root, fn)
                rp = f"{dst_dir}/{fn}"
                self.fs.put(lp, rp)

    def _list_local_ckpt_dirs(self) -> set[str]:
        """Return the set of local distributed checkpoint directory names."""
        if not os.path.isdir(self.src_root):
            return set()
        return {d for d in os.listdir(self.src_root)
                if os.path.isdir(os.path.join(self.src_root, d))}

    def _list_remote_top_dirs(self) -> set[str]:
        """Return the set of top-level checkpoint directory names currently on S3."""
        try:
            entries = self.fs.ls(self.s3_prefix, detail=True)
        except FileNotFoundError:
            return set()

        out = set()
        for e in entries:
            name = e.get("name", "")
            if name:
                base = name.rstrip("/").split("/")[-1]
                out.add(base)
        return out

    def _prune_remote_missing_locals(self):
        """
        Remove remote checkpoints that are no longer present locally.

        This enforces top-k consistency between local and remote stores by
        deleting any S3 checkpoint directories not found in the local root.
        """
        local_names = self._list_local_ckpt_dirs()
        remote_names = self._list_remote_top_dirs()
        to_delete = remote_names - local_names

        for d in to_delete:
            rp = f"{self.s3_prefix}/{d}"
            try:
                self.fs.rm(rp, recursive=True)
            except Exception as e:
                print(f"[MirrorDistCkptToS3] Warning: failed to delete remote '{rp}': {e}")

    # ---------------------------------------------------------------------
    # Core mirroring logic
    # ---------------------------------------------------------------------

    def _maybe_mirror(self, trainer):
        """
        Perform one mirror pass if a new finalized checkpoint exists.

        Only runs on rank 0 and skips checkpoints that are:
          - Not finalized yet,
          - Already mirrored,
          - Missing on disk (e.g., still being written).
        """
        if not self._is_global_zero(trainer):
            return

        last_path = getattr(self.cb, "last_model_path", None)
        if not last_path:
            return

        local_dir = str(ckpt_to_dir(last_path))
        if not os.path.isdir(local_dir):
            return

        if not self._is_finalized(local_dir):
            return

        if local_dir == self._last_mirrored_local_dir:
            return  # already synced

        name = os.path.basename(local_dir.rstrip("/"))
        remote_dir = f"{self.s3_prefix}/{name}"

        # Upload the new checkpoint directory
        self._upload_dir(local_dir, remote_dir)

        # Delete any remote checkpoints not present locally (top-k consistency)
        self._prune_remote_missing_locals()

        self._last_mirrored_local_dir = local_dir

    # ---------------------------------------------------------------------
    # Hooks: when to trigger mirroring
    # ---------------------------------------------------------------------
    def on_validation_end(self, trainer, pl_module):
        """Mirror checkpoints at the end of each validation phase."""
        self._maybe_mirror(trainer)

    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
        """
        Optionally mirror during training.
        For long runs, this can provide near-real-time S3 uploads.
        Comment out if you prefer to sync only after validation.
        """
        self._maybe_mirror(trainer)

    def on_train_end(self, trainer, pl_module):
        """Final mirror at the end of training."""
        self._maybe_mirror(trainer)

However, I wonder if there is a more native / natural way to do it with NeMo / BioNeMO:

• I've tried using S3CheckpointIO but kept running into serialization issues (even when setting use_weights_only=True wherever possible)
• I've tried passing the path directly to the ModelCheckpoint callback but in this case the s3 path is not read properly and I keep getting errors such as Invalid bucket name "s3:"