train: fallback to CPU on unsupported CUDA capability

Author: Dieg0Code

src/training/trainer_runtime.py

@@ -24,6 +24,14 @@ def resolve_trainer_hw() -> tuple[str, int, str]:
     requested_devices = cfg_int("trainer_devices")
     strategy = cfg_str("trainer_strategy")
     if torch.cuda.is_available():
+        capability = torch.cuda.get_device_capability(0)
+        if int(capability[0]) < 7:
+            log(
+                "Detected CUDA device with compute capability "
+                f"sm_{int(capability[0])}{int(capability[1])}, unsupported by current torch build. "
+                "Falling back to CPU.",
+            )
+            return "cpu", 1, "auto"
         available = max(1, torch.cuda.device_count())
         devices = min(requested_devices, available)
         if requested_devices > available:

tests/test_training_trainer_runtime.py

@@ -43,6 +43,7 @@ def test_resolve_trainer_hw_downgrades_ddp_spawn_when_one_gpu_available(self) ->
         CONFIG["trainer_strategy"] = "ddp_spawn"
         with (
             patch("training.trainer_runtime.torch.cuda.is_available", return_value=True),
+            patch("training.trainer_runtime.torch.cuda.get_device_capability", return_value=(8, 0)),
             patch("training.trainer_runtime.torch.cuda.device_count", return_value=1),
         ):
             accelerator, devices, strategy = resolve_trainer_hw()
@@ -55,13 +56,26 @@ def test_resolve_trainer_hw_keeps_ddp_spawn_when_two_gpus_available(self) -> Non
         CONFIG["trainer_strategy"] = "ddp_spawn"
         with (
             patch("training.trainer_runtime.torch.cuda.is_available", return_value=True),
+            patch("training.trainer_runtime.torch.cuda.get_device_capability", return_value=(8, 0)),
             patch("training.trainer_runtime.torch.cuda.device_count", return_value=2),
         ):
             accelerator, devices, strategy = resolve_trainer_hw()
         self.assertEqual(accelerator, "gpu")
         self.assertEqual(devices, 2)
         self.assertEqual(strategy, "ddp_spawn")
 
+    def test_resolve_trainer_hw_falls_back_to_cpu_for_unsupported_cuda_capability(self) -> None:
+        CONFIG["trainer_devices"] = 1
+        CONFIG["trainer_strategy"] = "auto"
+        with (
+            patch("training.trainer_runtime.torch.cuda.is_available", return_value=True),
+            patch("training.trainer_runtime.torch.cuda.get_device_capability", return_value=(6, 0)),
+        ):
+            accelerator, devices, strategy = resolve_trainer_hw()
+        self.assertEqual(accelerator, "cpu")
+        self.assertEqual(devices, 1)
+        self.assertEqual(strategy, "auto")
+
 
 if __name__ == "__main__":
     unittest.main()