|
24 | 24 |
|
25 | 25 | echo "Found CUDA installation at $CUDA_HOME"
|
26 | 26 |
|
27 |
| -cd /tmp |
| 27 | +# Detect architecture |
| 28 | +export TARGETARCH=${TARGETARCH:-$(uname -m)} |
| 29 | +if [ "${TARGETARCH}" = 'aarch64' ] || [ "${TARGETARCH}" = 'arm64' ]; then |
| 30 | + ARCH_PATH='sbsa' |
| 31 | +else |
| 32 | + ARCH_PATH='x86_64' |
| 33 | +fi |
28 | 34 |
|
29 |
| -# Download cuDNN archive |
30 |
| -CUDNN_ARCHIVE="cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz" |
31 |
| -echo "Looking for cuDNN archive: $CUDNN_ARCHIVE" |
| 35 | +echo "Architecture: ${ARCH_PATH}" |
32 | 36 |
|
33 |
| -if [ -f "$CUDNN_ARCHIVE" ]; then |
34 |
| - echo "Found cuDNN archive, installing from file..." |
| 37 | +cd /tmp |
35 | 38 |
|
36 |
| - # Extract cuDNN |
37 |
| - echo "Extracting cuDNN archive..." |
38 |
| - tar -xJf "$CUDNN_ARCHIVE" |
| 39 | +# Download cuDNN archive from NVIDIA |
| 40 | +CUDNN_ARCHIVE="cudnn-linux-${ARCH_PATH}-${CUDNN_VERSION}_cuda12-archive.tar.xz" |
| 41 | +CUDNN_URL="https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-${ARCH_PATH}/${CUDNN_ARCHIVE}" |
39 | 42 |
|
40 |
| - # Install cuDNN to CUDA directory |
41 |
| - echo "Installing cuDNN to $CUDA_HOME..." |
42 |
| - sudo cp cuda/include/cudnn*.h "$CUDA_HOME/include/" |
43 |
| - sudo cp cuda/lib64/libcudnn* "$CUDA_HOME/lib64/" |
44 |
| - sudo chmod a+r "$CUDA_HOME/include/cudnn*.h" "$CUDA_HOME/lib64/libcudnn*" |
| 43 | +echo "Downloading cuDNN from: $CUDNN_URL" |
| 44 | +echo "Archive name: $CUDNN_ARCHIVE" |
45 | 45 |
|
46 |
| - echo "cuDNN installed successfully from archive" |
47 |
| -else |
48 |
| - echo "cuDNN archive not found, using conda installation as fallback..." |
| 46 | +# Download cuDNN |
| 47 | +if ! wget -q "$CUDNN_URL" -O "$CUDNN_ARCHIVE"; then |
| 48 | + echo "ERROR: Failed to download cuDNN from $CUDNN_URL" |
| 49 | + exit 1 |
| 50 | +fi |
49 | 51 |
|
50 |
| - # Ensure conda environment is activated |
51 |
| - if [ -n "$CONDA_ENV" ]; then |
52 |
| - source /opt/miniconda3/etc/profile.d/conda.sh |
53 |
| - conda activate "$CONDA_ENV" |
54 |
| - fi |
| 52 | +echo "cuDNN download complete, extracting archive..." |
55 | 53 |
|
56 |
| - # Install cuDNN via conda |
57 |
| - conda install -c conda-forge cudnn="$CUDNN_VERSION" -y |
| 54 | +# Extract cuDNN |
| 55 | +if ! tar -xJf "$CUDNN_ARCHIVE"; then |
| 56 | + echo "ERROR: Failed to extract cuDNN archive" |
| 57 | + exit 1 |
| 58 | +fi |
58 | 59 |
|
59 |
| - echo "cuDNN installed successfully via conda" |
| 60 | +# Install cuDNN to CUDA directory |
| 61 | +echo "Installing cuDNN to $CUDA_HOME..." |
| 62 | +EXTRACTED_DIR="cudnn-linux-${ARCH_PATH}-${CUDNN_VERSION}_cuda12-archive" |
| 63 | + |
| 64 | +if [ -d "$EXTRACTED_DIR" ]; then |
| 65 | + sudo cp -a "$EXTRACTED_DIR/include/"* "$CUDA_HOME/include/" |
| 66 | + sudo cp -a "$EXTRACTED_DIR/lib/"* "$CUDA_HOME/lib64/" |
| 67 | + sudo chmod a+r "$CUDA_HOME/include/cudnn*.h" "$CUDA_HOME/lib64/libcudnn*" |
| 68 | + |
| 69 | + echo "cuDNN installed successfully" |
| 70 | +else |
| 71 | + echo "ERROR: Extracted directory not found: $EXTRACTED_DIR" |
| 72 | + exit 1 |
60 | 73 | fi
|
61 | 74 |
|
| 75 | +# Clean up downloaded files |
| 76 | +rm -f "$CUDNN_ARCHIVE" |
| 77 | +rm -rf "$EXTRACTED_DIR" |
| 78 | + |
62 | 79 | # Set cuDNN environment variables
|
63 | 80 | export LD_LIBRARY_PATH="$CUDA_HOME/lib64:$LD_LIBRARY_PATH"
|
64 | 81 |
|
|
0 commit comments