Fix numpy upgrade to 2.0.0 BUFSIZE import error (#5680)

Numpy 2.0.0 removed the BUFSIZE API, replacing it with 'numpy._core.umath.BUFSIZE'. The minimum numpy version supported by the new code is 1.26.1. If you need to support the lower numpy version, please tell me and we can add another path. Co-authored-by: Logan Adams <[email protected]>
microsoft · Jun 18, 2024 · 4000cee · 4000cee
1 parent a4cd550
commit 4000cee
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/deepspeed/autotuning/scheduler.py b/deepspeed/autotuning/scheduler.py
@@ -5,7 +5,11 @@
 
 import copy
 
-from numpy import BUFSIZE
+import numpy
+if numpy.__version__ < '2.0.0':
+    from numpy import BUFSIZE
+else:
+    from numpy._core.umath import BUFSIZE
 import json
 import subprocess
 import sys