Maintainance to pl_upgrades

environment.yml

@@ -15,10 +15,11 @@ dependencies:
   - biopython
   - numpy
   - pandas
-  - PyYAML==5.4.1
+  - PyYAML
+  - pytorch
   - requests
   - scipy
-  - tqdm==4.62.2
+  - tqdm
   - typing-extensions
   - wandb
   - modelcif==0.7
@@ -30,10 +31,8 @@ dependencies:
   - bioconda::hmmer
   - bioconda::hhsuite
   - bioconda::kalign2
-  - pytorch::pytorch=2.1
-  - pytorch::pytorch-cuda=12.1
   - pip:
-      - deepspeed==0.12.4
+      - deepspeed==0.14.5
       - dm-tree==0.1.6
       - git+https://github.com/NVIDIA/dllogger.git
       - flash-attn 

openfold/model/primitives.py

@@ -808,7 +808,7 @@ def _flash_attn(q, k, v, kv_mask):
     # [B_flat, N, 2 * H * C]
     kv = kv.reshape(*kv.shape[:-3], -1) 
     
-    kv_unpad, _, kv_cu_seqlens, kv_max_s = unpad_input(kv, kv_mask)
+    kv_unpad, _, kv_cu_seqlens, kv_max_s, _ = unpad_input(kv, kv_mask)
     kv_unpad = kv_unpad.reshape(-1, *kv_shape[-3:])
    
     out = flash_attn_varlen_kvpacked_func(

scripts/install_third_party_dependencies.sh

@@ -14,7 +14,7 @@ gunzip -c tests/test_data/sample_feats.pickle.gz > tests/test_data/sample_feats.
 python setup.py install
 
 echo "Download CUTLASS, required for Deepspeed Evoformer attention kernel"
-git clone https://github.com/NVIDIA/cutlass --depth 1
+git clone https://github.com/NVIDIA/cutlass --branch v3.6.0 --depth 1
 conda env config vars set CUTLASS_PATH=$PWD/cutlass
 
 # This setting is used to fix a worker assignment issue during data loading