Updated DeepSpeed version, added inplace_safe for template offloading

environment.yml

@@ -30,7 +30,7 @@ dependencies:
   - bioconda::kalign2==2.04
   - pytorch::pytorch=1.12.*
   - pip:
-      - git+https://github.com/microsoft/DeepSpeed.git@4388a60  # Replace when version becomes available
+      - deepspeed==0.12.4
       - dm-tree==0.1.6
       - git+https://github.com/NVIDIA/dllogger.git
       - git+https://github.com/Dao-AILab/flash-attention.git@5b838a8

openfold/model/template.py

@@ -474,6 +474,7 @@ def embed_templates_offload(
             chunk_size=model.globals.chunk_size,
             use_deepspeed_evo_attention=model.globals.use_deepspeed_evo_attention,
             use_lma=model.globals.use_lma,
+            inplace_safe=inplace_safe,
             _mask_trans=model.config._mask_trans,
         )
 
@@ -592,6 +593,7 @@ def embed_templates_average(
             chunk_size=model.globals.chunk_size,
             use_deepspeed_evo_attention=model.globals.use_deepspeed_evo_attention,
             use_lma=model.globals.use_lma,
+            inplace_safe=inplace_safe,
             _mask_trans=model.config._mask_trans,
         )