[devops] fix extention building (#5427)

examples/language/llama2/README.md

@@ -53,7 +53,7 @@ We follow the hyperparameter settings from the original LLaMA paper. We use Adam
 Please install the latest ColossalAI from source.
 
 ```bash
-CUDA_EXT=1 pip install -U git+https://github.com/hpcaitech/ColossalAI
+BUILD_EXT=1 pip install -U git+https://github.com/hpcaitech/ColossalAI
 ```
 
 Then install other dependencies.

examples/language/openmoe/README.md

@@ -17,7 +17,7 @@
 Please install the latest ColossalAI from source.
 
 ```bash
-CUDA_EXT=1 pip install -U git+https://github.com/hpcaitech/ColossalAI
+BUILD_EXT=1 pip install -U git+https://github.com/hpcaitech/ColossalAI
 ```
 
 Then install dependencies.

extensions/utils.py

@@ -154,7 +154,7 @@ def check_cuda_availability():
 def set_cuda_arch_list(cuda_dir):
     """
     This function sets the PyTorch TORCH_CUDA_ARCH_LIST variable for ahead-of-time extension compilation.
-    Ahead-of-time compilation occurs when CUDA_EXT=1 is set when running 'pip install'.
+    Ahead-of-time compilation occurs when BUILD_EXT=1 is set when running 'pip install'.
     """
     cuda_available = check_cuda_availability()
 

setup.py

@@ -70,7 +70,7 @@ def get_version() -> str:
 if BUILD_EXT:
     if not TORCH_AVAILABLE:
         raise ModuleNotFoundError(
-            "[extension] PyTorch is not found while CUDA_EXT=1. You need to install PyTorch first in order to build CUDA extensions"
+            "[extension] PyTorch is not found while BUILD_EXT=1. You need to install PyTorch first in order to build CUDA extensions"
         )
 
     from extensions import ALL_EXTENSIONS