From af4cc28de2b6dfb4d234fd41715fc8c256b55df8 Mon Sep 17 00:00:00 2001
From: hailin <hailin@gdzx.xyz>
Date: Fri, 8 Aug 2025 21:43:40 +0800
Subject: [PATCH] .

---
 .../install_deepspeed_src_mamba_single_gpu.sh | 41 +++++++++++--------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/dspeed/install_deepspeed_src_mamba_single_gpu.sh b/dspeed/install_deepspeed_src_mamba_single_gpu.sh
index f9772e7..476a951 100755
--- a/dspeed/install_deepspeed_src_mamba_single_gpu.sh
+++ b/dspeed/install_deepspeed_src_mamba_single_gpu.sh
@@ -173,32 +173,37 @@ log "Build & install DeepSpeed"
 cd "${REPO_PATH}"
 pip uninstall -y deepspeed >/dev/null 2>&1 || true
 rm -rf build
+
+# —— 编译前置（强烈建议）——
+python - <<'PY'
+import sys, subprocess
+subprocess.check_call([sys.executable, "-m", "pip", "install", "-U", "pip", "setuptools", "wheel", "ninja"])
+PY
+
+
+# 并行编译提速（按 CPU 核心数调）
+export MAX_JOBS=${MAX_JOBS:-$(nproc)}
+
+# —— DeepSpeed 功能开关 —— 
+# 编译核心 CUDA/C++ 扩展（必须）
 export DS_BUILD_OPS=1
-export DS_BUILD_AIO=1
-export DS_BUILD_FUSED_ADAM=1
-export DS_BUILD_CPU_ADAM=1
-export DS_BUILD_EVOFORMER_ATTN=0
 
-export DS_BUILD_AIO=1
+# 训练常用加速内核
+export DS_BUILD_TRANSFORMER=1
+export DS_BUILD_SPARSE_ATTN=1
+export DS_BUILD_FLASH_ATTN=1
+
+# 常用优化器（强烈建议开）
 export DS_BUILD_FUSED_ADAM=1
 export DS_BUILD_CPU_ADAM=1
 
-# 全部关掉和推理/transformer/CUTLASS/flash/sparse 有关的东西
-export DS_BUILD_EVOFORMER_ATTN=0
-export DS_BUILD_INFERENCE=0
-export DS_BUILD_TRANSFORMER=0
-export DS_BUILD_SPARSE_ATTN=0
-export DS_BUILD_SDP_ATTN=0
-export DS_BUILD_FLASH_ATTN=0
-export DS_BUILD_QKV_GEMM=0
-export DS_BUILD_CUTLASS=0
-unset CUTLASS_PATH
+# 可选：异步 IO（Linux 才生效；Windows 不支持 AIO/GDS）
+export DS_BUILD_AIO=1
 
-# 如需推理内核：export DS_BUILD_TRANSFORMER=1
-
-# 提供更详细日志便于排错
+# 安装（带详细日志）
 pip install -v .
 
+
 # =============================
 # 7. 验证安装
 # =============================