yindo/flash-attention-prebuild-wheels-rocm
1
0
Fork 0
mirror of https://github.com/BillyOutlast/flash-attention-prebuild-wheels-rocm.git synced 2026-07-01 01:27:54 -04:00
Code Issues Packages Projects Releases 2 Wiki Activity

fix: Update build parallel

Browse Source
This commit is contained in:
Junya Morioka
2025-12-05 00:47:29 +09:00
parent 612ab6d2ca
commit cdfbcd8034
2 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/test-self-hosted-linux.yml
+1 -1
View File
@@ -89,7 +89,7 @@ jobs:
chmod +x build_linux.sh chmod +x build_linux.sh
./build_linux.sh ${{ matrix.flash-attn-version }} ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }} ./build_linux.sh ${{ matrix.flash-attn-version }} ${{ matrix.python-version }} ${{ matrix.torch-version }} ${{ matrix.cuda-version }}
wheel_name=$(basename $(ls flash-attention/dist/*.whl | head -n 1)) wheel_name=$(basename $(ls flash-attention/dist/*.whl | head -n 1))
echo "wheel_name=$wheel_name" >> $GITHUB_ENV echo "WHEEL_NAME=$wheel_name" >> $GITHUB_OUTPUT
- name: Install Test - name: Install Test
shell: bash shell: bash
build_linux.sh
+12 -6
View File
@@ -54,15 +54,21 @@ echo " RAM: ${RAM_GB}GB"
# Determine MAX_JOBS and NVCC_THREADS based on system resources # Determine MAX_JOBS and NVCC_THREADS based on system resources
if [[ -z "${MAX_JOBS:-}" && -z "${NVCC_THREADS:-}" ]]; then if [[ -z "${MAX_JOBS:-}" && -z "${NVCC_THREADS:-}" ]]; then
# Calculate max product based on following constraints: # Calculate max product based on following constraints:
# - MAX_JOBS x NVCC_THREADS <= NUM_THREADS # - MAX_JOBS x NVCC_THREADS(<= 4) <= NUM_THREADS
# - 2.5GB x MAX_JOBS x NVCC_THREADS <= RAM_GB # - 2.5GB x MAX_JOBS x NVCC_THREADS(<= 4) <= RAM_GB
MAX_PRODUCT_CPU=$NUM_THREADS MAX_PRODUCT_CPU=$NUM_THREADS
MAX_PRODUCT_RAM=$(awk -v ram="$RAM_GB" 'BEGIN {print int(ram / 2.5)}') MAX_PRODUCT_RAM=$(awk -v ram="$RAM_GB" 'BEGIN {print int(ram / 2.25)}')
MAX_PRODUCT=$((MAX_PRODUCT_CPU < MAX_PRODUCT_RAM ? MAX_PRODUCT_CPU : MAX_PRODUCT_RAM)) MAX_PRODUCT=$((MAX_PRODUCT_CPU < MAX_PRODUCT_RAM ? MAX_PRODUCT_CPU : MAX_PRODUCT_RAM))
# Set MAX_JOBS = NVCC_THREADS = floor(sqrt(MAX_PRODUCT)) # Set MAX_JOBS and NVCC_THREADS so that MAX_JOBS x NVCC_THREADS ≈ MAX_PRODUCT with NVCC_THREADS <= 4
MAX_JOBS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}') BASE_THREADS=$(awk -v max="$MAX_PRODUCT" 'BEGIN {print int(sqrt(max))}')
NVCC_THREADS=$MAX_JOBS if (( BASE_THREADS <= 4 )); then
MAX_JOBS=$BASE_THREADS
NVCC_THREADS=$BASE_THREADS
else
NVCC_THREADS=4
MAX_JOBS=$((MAX_PRODUCT / NVCC_THREADS))
fi
# Ensure minimum values of 1 # Ensure minimum values of 1
MAX_JOBS=$((MAX_JOBS < 1 ? 1 : MAX_JOBS)) MAX_JOBS=$((MAX_JOBS < 1 ? 1 : MAX_JOBS))