From 4a65d60bd547c6675490eee68686e97292b59a7d Mon Sep 17 00:00:00 2001 From: Jimin Ha Date: Tue, 10 Sep 2024 23:28:28 -0700 Subject: [PATCH 1/3] Update esmfold model not to use param_buffer_assignment --- examples/protein-folding/run_esmfold.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/examples/protein-folding/run_esmfold.py b/examples/protein-folding/run_esmfold.py index 4fa8d7a44..b924d611a 100644 --- a/examples/protein-folding/run_esmfold.py +++ b/examples/protein-folding/run_esmfold.py @@ -82,6 +82,9 @@ def convert_outputs_to_pdb(outputs): test_protein = "MGAGASAEEKHSRELEKKLKEDAEKDARTVKLLLLGAGESGKSTIVKQMKIIHQDGYSLEECLEFIAIIYGNTLQSILAIVRAMTTLNIQYGDSARQDDARKLMHMADTIEEGTMPKEMSDIIQRLWKDSGIQACFERASEYQLNDSAGYYLSDLERLVTPGYVPTEQDVLRSRVKTTGIIETQFSFKDLNFRMFDVGGQRSERKKWIHCFEGVTCIIFIAALSAYDMVLVEDDEVNRMHESLHLFNSICNHRYFATTSIVLFLNKKDVFFEKIKKAHLSICFPDYDGPNTYEDAGNYIKVQFLELNMRRDVKEIYSHMTCATDTQNVKFVFDAVTDIIIKENLKDCGLF" # len = 350 tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1") +# Set _supports_param_buffer_assignment to False since facebook/esmfold_v1's encoder weights are float16. +# Without this fix, we will get the weights to be loaded with float16 on gaudi2,gaudi3 and runtime error on gaudi1 +EsmForProteinFolding._supports_param_buffer_assignment = False model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=False) model = model.to(device) From 19763afa9f98d23cf2fdbd9af04337eb45fc8602 Mon Sep 17 00:00:00 2001 From: Jimin Ha Date: Tue, 10 Sep 2024 23:49:03 -0700 Subject: [PATCH 2/3] comment fix --- examples/protein-folding/run_esmfold.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/protein-folding/run_esmfold.py b/examples/protein-folding/run_esmfold.py index b924d611a..4fe7359c5 100644 --- a/examples/protein-folding/run_esmfold.py +++ b/examples/protein-folding/run_esmfold.py @@ -83,9 +83,10 @@ def convert_outputs_to_pdb(outputs): tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1") # Set _supports_param_buffer_assignment to False since facebook/esmfold_v1's encoder weights are float16. -# Without this fix, we will get the weights to be loaded with float16 on gaudi2,gaudi3 and runtime error on gaudi1 +# Without this fix, we will have the weights loaded with float16 on gaudi2,gaudi3 and runtime error on gaudi1 EsmForProteinFolding._supports_param_buffer_assignment = False -model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=False) +#model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=False) +model = EsmForProteinFolding.from_pretrained("/root/tf/models--facebook--esmfold_v1/snapshots/75a3841ee059df2bf4d56688166c8fb459ddd97a/", low_cpu_mem_usage=False) model = model.to(device) # Uncomment this line if you're folding longer (over 600 or so) sequences From 4189d9cdd5e8d0e0e865cbe6f2dce81672bc8748 Mon Sep 17 00:00:00 2001 From: Jimin Ha Date: Tue, 10 Sep 2024 23:55:58 -0700 Subject: [PATCH 3/3] Fix changes error --- examples/protein-folding/run_esmfold.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/protein-folding/run_esmfold.py b/examples/protein-folding/run_esmfold.py index 4fe7359c5..6941e6e5c 100644 --- a/examples/protein-folding/run_esmfold.py +++ b/examples/protein-folding/run_esmfold.py @@ -85,8 +85,7 @@ def convert_outputs_to_pdb(outputs): # Set _supports_param_buffer_assignment to False since facebook/esmfold_v1's encoder weights are float16. # Without this fix, we will have the weights loaded with float16 on gaudi2,gaudi3 and runtime error on gaudi1 EsmForProteinFolding._supports_param_buffer_assignment = False -#model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=False) -model = EsmForProteinFolding.from_pretrained("/root/tf/models--facebook--esmfold_v1/snapshots/75a3841ee059df2bf4d56688166c8fb459ddd97a/", low_cpu_mem_usage=False) +model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=False) model = model.to(device) # Uncomment this line if you're folding longer (over 600 or so) sequences