Help with Esm-2 model

[Afganitia]

Hello,

I am trying to use the big Esm-2 models to create embeddings. I first tried the 15b, and then the 3b. I got OoM errors all the time. The 15b wouldn't even load to the gpus (never reached step2)! I tried to use multigpu in one node, checkpoint activation and cpu sharding, but not even the 3b was capable of doing a single sequence (around 300 aminoacid length).

I am using modest hardware, with 128 cores, 1tb ram and 4xNvidia A100 80gb vram.

So, question. How much hardware is required to run the 3b model?

Code

import os
import socket
import torch
import torch.distributed as dist
import esm
import pandas as pd
import logging
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.utils.checkpoint import checkpoint

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger()

def check_port_availability(port):
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        return s.connect_ex(('localhost', port)) != 0

def check_gpu_availability():
    return torch.cuda.device_count()

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)
    [logger.info](http://logger.info/)(f"Process {rank} set up")

def cleanup():
    dist.destroy_process_group()

def main(rank, world_size):
    setup(rank, world_size)
    print('step1')
    # Load ESM-2 model and move it to the GPU
    # with fsdp
    try:
        model, alphabet = esm.pretrained.esm2_t48_15B_UR50D()
        model = [model.to](http://model.to/)(rank)
        model = FSDP(model, device_id=rank)  # Wrap model with FSDP
        batch_converter = alphabet.get_batch_converter()
        model.eval()  # Disables dropout for deterministic results
    except Exception as e:
        logger.error(f"Error loading model: {e}")
        cleanup()
        return

    print('step2')

    file_path = 'path/to/file.csv'
    try:
        df = pd.read_csv(file_path)
        data2 = list(df[['Uniprot_id', 'Secuencias']].itertuples(index=False, name=None))
    except Exception as e:
        logger.error(f"Error loading CSV file: {e}")
        cleanup()
        return

    try:
        data = data2[:10] # choosing only 10 sequences for debugging
        batch_labels, batch_strs, batch_tokens = batch_converter(data)
        batch_tokens = [batch_tokens.to](http://batch_tokens.to/)(rank)  # Move tokens to GPU
        batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
    except Exception as e:
        logger.error(f"Error preparing data: {e}")
        cleanup()
        return

    print("Prep work done")

    # checkpointing
    def forward_func(inputs):
        return model(inputs, repr_layers=[48], return_contacts=True)

    # Extract per-residue representations (on GPU)
    try:
        with torch.no_grad():
            results = checkpoint(forward_func, batch_tokens)
        token_representations = results["representations"][48]
    except Exception as e:
        logger.error(f"Error extracting representations: {e}")
        cleanup()
        return

    # Generate per-sequence representations via averaging
    sequence_representations = []
    try:
        for i, tokens_len in enumerate(batch_lens):
            sequence_representations.append(token_representations[i, 1: tokens_len - 1].mean(0))
    except Exception as e:
        logger.error(f"Error generating sequence representations: {e}")
        cleanup()
        return

    # Saving
    if rank == 0:
        try:
            with open('sequence_embeddings_esm.txt', 'w') as f:
                for label, representation in zip(batch_labels, sequence_representations):
                    f.write(f"{label}\t{' '.join(map(str, representation.cpu().numpy()))}\n")
            [logger.info](http://logger.info/)("Saved sequence representations to 'sequence_embeddings_esm.txt'")
        except Exception as e:
            logger.error(f"Error saving representations: {e}")

    cleanup()

if __name__ == "__main__":
    # Check GPU availability
    world_size = check_gpu_availability()
    if world_size < 1:
        logger.error("No GPUs available. Exiting.")
        exit(1)

    # Parse the rank and world size from environment variables set by torchrun
    rank = int(os.environ['RANK'])
    world_size = int(os.environ['WORLD_SIZE'])

    # Run the main script
    main(rank, world_size)

# code run using comand:
#torchrun --nproc_per_node=4 esm_33b_whatever.py
  

[lzygitk7]

I use A40 46gb GRAM, and 3b is good for me, a test around 1800 aa length would reach the cell to get an OoM error.

[Afganitia]

As I suspected, I believe I am doing something wrong. Thanks.