enforce fp32 and torch kernels for triangle attention and multiplicative on V100

Author: zhangyuxuann

runner/batch_inference.py

@@ -36,7 +36,12 @@
 from protenix.data.json_parser import lig_file_to_atom_info
 from protenix.data.utils import pdb_to_cif
 from protenix.utils.logger import get_logger
-from runner.inference import InferenceRunner, download_infercence_cache, infer_predict
+from runner.inference import (
+    InferenceRunner,
+    download_infercence_cache,
+    infer_predict,
+    update_gpu_compatible_configs,
+)
 from runner.msa_search import msa_search, update_infer_json
 
 logger = get_logger(__name__)
@@ -166,6 +171,7 @@ def get_default_runner(
     n_cycle: int = 10,
     n_step: int = 200,
     n_sample: int = 5,
+    dtype: str = "bf16",
     model_name: str = "protenix_base_default_v0.5.0",
     use_msa: bool = True,
     trimul_kernel="cuequivariance",
@@ -184,23 +190,28 @@ def get_default_runner(
         configs.seeds = seeds
     model_name = configs.model_name
     _, model_size, model_feature, model_version = model_name.split("_")
-    logger.info(
-        f"Inference by Protenix: model_size: {model_size}, with_feature: {model_feature.replace('-', ',')}, model_version: {model_version}"
-    )
     model_specfics_configs = ConfigDict(model_configs[model_name])
     # update model specific configs
     configs.update(model_specfics_configs)
     # the user input configs has the highest priority
     configs.model.N_cycle = n_cycle
     configs.sample_diffusion.N_sample = n_sample
     configs.sample_diffusion.N_step = n_step
+    configs.dtype = dtype
     configs.use_msa = use_msa
     configs.triangle_multiplicative = trimul_kernel
     configs.triangle_attention = triatt_kernel
     configs.enable_diffusion_shared_vars_cache = enable_cache
     configs.enable_efficient_fusion = enable_fusion
     configs.enable_tf32 = enable_tf32
 
+    configs = update_gpu_compatible_configs(configs)
+    logger.info(
+        f"Inference by Protenix: model_size: {model_size}, with_feature: {model_feature.replace('-', ',')}, model_version: {model_version}, dtype: {configs.dtype}"
+    )
+    logger.info(
+        f"Triangle_multiplicative kernel: {trimul_kernel}, Triangle_attention kernel: {triatt_kernel}"
+    )
     logger.info(
         f"enable_diffusion_shared_vars_cache: {configs.enable_diffusion_shared_vars_cache}, "
         + f"enable_efficient_fusion: {configs.enable_efficient_fusion}, enable_tf32: {configs.enable_tf32}"
@@ -217,6 +228,7 @@ def inference_jsons(
     n_cycle: int = 10,
     n_step: int = 200,
     n_sample: int = 5,
+    dtype: str = "bf16",
     model_name: str = "protenix_base_default_v0.5.0",
     trimul_kernel="cuequivariance",
     triatt_kernel="triattention",
@@ -255,6 +267,7 @@ def inference_jsons(
         n_cycle,
         n_step,
         n_sample,
+        dtype,
         model_name,
         use_msa,
         trimul_kernel,
@@ -290,6 +303,7 @@ def protenix_cli():
 @click.option("-c", "--cycle", type=int, default=10, help="pairformer cycle number")
 @click.option("-p", "--step", type=int, default=200, help="diffusion step")
 @click.option("-e", "--sample", type=int, default=5, help="sample number")
+@click.option("-d", "--dtype", type=str, default="bf16", help="sample number")
 @click.option(
     "-n",
     "--model_name",
@@ -349,6 +363,7 @@ def predict(
     cycle,
     step,
     sample,
+    dtype,
     model_name,
     use_msa,
     use_default_params,
@@ -403,9 +418,6 @@ def predict(
         "deepspeed",
         "torch",
     ], "Kernel to use for triangle attention. Options: 'triattention', 'cuequivariance', 'deepspeed', 'torch'."
-    logger.info(
-        f"Triangle_multiplicative kernel: {trimul_kernel}, Triangle_attention kernel: {triatt_kernel}"
-    )
     seeds = list(map(int, seeds.split(",")))
     inference_jsons(
         input,
@@ -415,6 +427,7 @@ def predict(
         n_cycle=cycle,
         n_step=step,
         n_sample=sample,
+        dtype=dtype,
         model_name=model_name,
         trimul_kernel=trimul_kernel,
         triatt_kernel=triatt_kernel,

runner/inference.py

@@ -372,6 +372,30 @@ def main(configs: Any) -> None:
     infer_predict(runner, configs)
 
 
+def update_gpu_compatible_configs(configs: Any) -> None:
+    def is_gpu_capability_between_7_and_8():
+        # 7.0 <= device_capability < 8.0
+        if not torch.cuda.is_available():
+            return False
+
+        capability = torch.cuda.get_device_capability()
+        major, minor = capability
+        cc = major + minor / 10.0
+        if 7.0 <= cc < 8.0:
+            return True
+        return False
+
+    if is_gpu_capability_between_7_and_8():
+        # Some kernels and BF16 aren’t supported on V100 — enforce specific configurations to work around it.
+        configs.dtype = "fp32"
+        configs.triangle_attention = "torch"
+        configs.triangle_multiplicative = "torch"
+        logger.info(
+            "GPU capability is between 7.0 and 8.0, enforce fp32 and torch kernels for triangle attention and multiplicative."
+        )
+    return configs
+
+
 def run() -> None:
     LOG_FORMAT = "%(asctime)s,%(msecs)-3d %(levelname)-8s [%(filename)s:%(lineno)s %(funcName)s] %(message)s"
     logging.basicConfig(
@@ -390,11 +414,12 @@ def run() -> None:
     model_name = configs.model_name
     _, model_size, model_feature, model_version = model_name.split("_")
     logger.info(
-        f"Inference by Protenix: model_size: {model_size}, with_feature: {model_feature.replace('-',', ')}, model_version: {model_version}"
+        f"Inference by Protenix: model_size: {model_size}, with_feature: {model_feature.replace('-',', ')}, model_version: {model_version}, dtype: {configs.dtype}"
     )
     model_specfics_configs = ConfigDict(model_configs[model_name])
     # update model specific configs
     configs.update(model_specfics_configs)
+    configs = update_gpu_compatible_configs(configs)
     logger.info(
         f"Triangle_multiplicative kernel: {configs.triangle_multiplicative}, Triangle_attention kernel: {configs.triangle_attention}"
     )

setup.py

@@ -38,7 +38,7 @@
 setup(
     name="protenix",
     python_requires=">=3.11",
-    version="0.7.0",
+    version="0.7.1",
     description="A trainable PyTorch reproduction of AlphaFold 3.",
     long_description=long_description,
     long_description_content_type="text/markdown",