From a08cffd63ec68b980b45c132f556d10a0029e548 Mon Sep 17 00:00:00 2001
From: Augustin Zidek <augustinzidek@google.com>
Date: Mon, 18 Nov 2024 16:41:32 +0000
Subject: [PATCH] Fail hard if running on a device with GPU capability < 8.0

PiperOrigin-RevId: 697640745
---
 run_alphafold.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/run_alphafold.py b/run_alphafold.py
index 4681134..4480afc 100644
--- a/run_alphafold.py
+++ b/run_alphafold.py
@@ -610,6 +610,17 @@ def main(_):
     print(f'Failed to create output directory {_OUTPUT_DIR.value}: {e}')
     raise
 
+  if _RUN_INFERENCE.value:
+    # Fail early on incompatible devices, but only if we're running inference.
+    gpu_devices = jax.local_devices(backend='gpu')
+    if gpu_devices and float(gpu_devices[0].compute_capability) < 8.0:
+      raise ValueError(
+          'There are currently known unresolved numerical issues with using'
+          ' devices with compute capability less than 8.0. See '
+          ' https://github.com/google-deepmind/alphafold3/issues/59 for'
+          ' tracking.'
+      )
+
   notice = textwrap.wrap(
       'Running AlphaFold 3. Please note that standard AlphaFold 3 model'
       ' parameters are only available under terms of use provided at'
-- 
GitLab