fix check CUDA_DEVICE_MAX_CONNECTIONS

InternLM · Dec 2, 2024 · 317f18c · 317f18c
1 parent 4a6b453
commit 317f18c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/internlm/utils/common.py b/internlm/utils/common.py
@@ -250,8 +250,9 @@ def enable_pytorch_expandable_segments():
 
 
 def check_cuda_env():
-    if os.getenv("CUDA_DEVICE_MAX_CONNECTIONS") is None:
-        logger.warning("Env var CUDA_DEVICE_MAX_CONNECTIONS has not be set, please note this!")
+    max_connections = os.getenv("CUDA_DEVICE_MAX_CONNECTIONS")
+    assert max_connections is not None, "Env var CUDA_DEVICE_MAX_CONNECTIONS has not been set, please set it to 1!"
+    assert max_connections == '1', "Env var CUDA_DEVICE_MAX_CONNECTIONS is set to {}, but it should be set to 1!".format(max_connections)
 
 
 class DummyProfile: