diff --git a/evo2/configs/evo2-7b-1m.yml b/evo2/configs/evo2-7b-1m.yml index 5ee5461..4fc408f 100644 --- a/evo2/configs/evo2-7b-1m.yml +++ b/evo2/configs/evo2-7b-1m.yml @@ -49,7 +49,7 @@ mha_out_proj_bias: True hyena_out_proj_bias: True hyena_flip_x1x2: False qkv_proj_bias: False -use_fp8_input_projections: True +use_fp8_input_projections: False max_seqlen: 1048576 max_batch_size: 1 final_norm: True diff --git a/evo2/test/test_evo2_generation.py b/evo2/test/test_evo2_generation.py index bc420b7..928854f 100644 --- a/evo2/test/test_evo2_generation.py +++ b/evo2/test/test_evo2_generation.py @@ -4,6 +4,7 @@ from importlib import resources from pathlib import Path from typing import List, Optional, Union import numpy as np +import time import torch @@ -65,6 +66,8 @@ def generate_and_score(*, sequences, model, generations_per_prompt=5, n_tokens=5 target = targets[i] with torch.inference_mode(): + if torch.cuda.is_available(): torch.cuda.synchronize() + elapsed_time = -time.perf_counter() generated = model.generate( prompt_seqs=[prompt], n_tokens=n_tokens, @@ -72,6 +75,9 @@ def generate_and_score(*, sequences, model, generations_per_prompt=5, n_tokens=5 top_k=top_k, top_p=top_p, ) + if torch.cuda.is_available(): torch.cuda.synchronize() + elapsed_time += time.perf_counter() + print(f"[{i}] Time for model.generate: {elapsed_time:.3f} s") decoded_seq = generated.sequences[0] # Assuming generate returns list of sequences score = calculate_sequence_identity(decoded_seq, target) @@ -95,7 +101,7 @@ def main(): parser = argparse.ArgumentParser(description="Test Evo2 Model Generation") parser.add_argument("--model_name", choices=['evo2_7b', 'evo2_40b', 'evo2_1b_base', 'evo2_20b'], default='evo2_7b', help="Model to test (supports evo2_7b, evo2_40b, evo2_1b_base, evo2_20b)") - + parser.add_argument("--local_path", type=str, default=None) args = parser.parse_args() # Reduce CUDA memory fragmentation for large models (e.g. evo2_20b) @@ -105,7 +111,7 @@ def main(): torch.manual_seed(1) torch.cuda.manual_seed(1) - model = Evo2(args.model_name) + model = Evo2(args.model_name, local_path=args.local_path) # Test parameters: greedy sampling of 500 tokens test_params = { @@ -145,4 +151,4 @@ def main(): print(f"\nTest Failed: Expected {expected_score}%, got {mean_score}%") if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/pyproject.toml b/pyproject.toml index 4347bbc..77b3f6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "evo2" version = "0.5.3" description = "Genome modeling across all domains of life" readme = "README.md" -requires-python = ">=3.11,<3.13" +requires-python = ">=3.10,<3.13" license = {file = "LICENSE"} authors = [ {name = "Evo 2 Team"},