fix: condaenv package dependency

Author: Jensen246

rdagent/scenarios/finetune/conda/llm_finetune_requirements.txt

@@ -2,8 +2,9 @@
 # Equivalent to: rdagent/scenarios/finetune/docker/llm_finetune_docker/Dockerfile
 # Docker base: hiyouga/llamafactory:0.9.4 uses PyTorch 2.6.0 + CUDA 12.4 + flash-attn 2.7.4
 
-# PyTorch 2.6.0 with CUDA 12.4 (matches Docker base image)
---index-url https://download.pytorch.org/whl/cu124
+# PyTorch 2.6.0 with CUDA 12.8 (for B200 GPUs)
+# Note: Change to cu124 for CUDA 12.4 machines
+--index-url https://download.pytorch.org/whl/cu128
 torch==2.6.0
 torchvision==0.21.0
 
@@ -13,8 +14,8 @@ torchvision==0.21.0
 # Core LlamaFactory package (PyPI latest is 0.9.3, Docker uses 0.9.4 from GitHub)
 llamafactory==0.9.3
 
-# FlashAttention-2 (matches Docker base image version)
-flash-attn==2.7.4
+# FlashAttention-2: installed separately via llm_finetune_flash_attn.txt
+# (requires torch installed first, and --no-build-isolation flag)
 
 # Additional dependencies (matches Dockerfile line 17)
 bitsandbytes>=0.39.0

rdagent/utils/env.py

@@ -835,10 +835,25 @@ class FTCondaEnv(LocalEnv[FTCondaConf]):
 
     def prepare(self) -> None:
         try:
+            env_name = self.conf.conda_env_name
+
+            # Skip if already prepared
+            if env_name in _CONDA_ENV_PREPARED:
+                return
+
+            # Step 1: Install base dependencies (torch, llamafactory, etc.)
             req_file = FT_CONDA_CONFIG_DIR / "llm_finetune_requirements.txt"
-            _prepare_conda_env(self.conf.conda_env_name, req_file)
+            _prepare_conda_env(env_name, req_file)
+
+            # Step 2: Install flash-attn (requires torch first, uses --no-build-isolation)
+            # --no-cache-dir: avoid cross-filesystem hardlink error when /tmp and ~/.cache/pip are on different mounts
+            print("[yellow]Installing flash-attn (compiling, may take a few minutes)...[/yellow]")
+            subprocess.check_call(
+                f"conda run -n {env_name} pip install 'flash-attn>=2.5.6,<=2.7.4' --no-build-isolation --no-cache-dir",
+                shell=True,
+            )
+
             # Re-update bin_path after prepare() in case the conda env was just created
-            # This fixes the issue where bin_path is empty if queried before env exists
             if not self.conf.bin_path:
                 self.conf._update_bin_path()
         except Exception as e: