zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Use the correct repo for arm64 CUDA

Browse Source
This commit is contained in:
Cheng
2025-11-18 09:34:12 +09:00
parent 82faffa7bb
commit 0f8c8c8a11
2 changed files with 11 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
.github/actions/build-linux/action.yml vendored
View File

@@ -9,9 +9,10 @@ inputs:
runs: runs:
using: "composite" using: "composite"
steps: steps:
- name: Set cmake args - name: Install Python package
shell: bash shell: sh
env: env:
DEBUG: 1
CMAKE_ARGS_CPU: >- CMAKE_ARGS_CPU: >-
-DCMAKE_COMPILE_WARNING_AS_ERROR=ON -DCMAKE_COMPILE_WARNING_AS_ERROR=ON
-DMLX_BUILD_CUDA=OFF -DMLX_BUILD_CUDA=OFF
@@ -20,22 +21,16 @@ runs:
-DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc -DCMAKE_CUDA_COMPILER=/usr/local/${{ inputs.toolkit }}/bin/nvcc
-DMLX_BUILD_CUDA=ON -DMLX_BUILD_CUDA=ON
run: | run: |
CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}" export CMAKE_ARGS="$CMAKE_ARGS_${{ startsWith(inputs.toolkit, 'cuda') && 'CUDA' || 'CPU' }}"
# FIXME: Unable to build tests without libnccl on arm64.
if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then if ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64' }} ; then
# Can not build tests when the built executables can not run.
CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF" CMAKE_ARGS="$CMAKE_ARGS -DMLX_BUILD_TESTS=OFF"
fi fi
pip install --no-build-isolation -e ".[dev]" -v
# Pass the CMAKE_ARGS to following steps.
echo CMAKE_ARGS="$CMAKE_ARGS" > $GITHUB_ENV echo CMAKE_ARGS="$CMAKE_ARGS" > $GITHUB_ENV
- name: Install Python package
shell: sh
env:
DEBUG: 1
run: pip install --no-build-isolation -e ".[dev]" -v
- name: Generate package stubs - name: Generate package stubs
# FIXME: Unable to generate stubs without libnccl on arm64.
if: ${{ !(startsWith(inputs.toolkit, 'cuda') && runner.arch == 'arm64') }}
shell: sh shell: sh
run: | run: |
pip install typing_extensions pip install typing_extensions

12
.github/actions/setup-linux/action.yml vendored
View File

@@ -60,17 +60,13 @@ runs:
"cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0" "cuda-13.0": "libcudnn9-dev-cuda-13 cuda-toolkit-13-0"
} }
run: | run: |
export ARCH=${{ runner.arch == 'arm64' && 'arm64' || 'x86_64' }} export ARCH=${{ runner.arch == 'arm64' && 'sbsa' || 'x86_64' }}
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/$ARCH/cuda-keyring_1.1-1_all.deb
sudo dpkg -i cuda-keyring_1.1-1_all.deb sudo dpkg -i cuda-keyring_1.1-1_all.deb
sudo apt-get update sudo apt-get update
sudo apt-get install -y ${{ fromJson(env.PACKAGES)[inputs.toolkit] }} sudo apt-get install -y \
libnccl2 libnccl-dev \
- name: Install NCCL ${{ fromJson(env.PACKAGES)[inputs.toolkit] }}
# For some reason NVIDIA did not provide arm64 package for libnccl2.
if: ${{ startsWith(inputs.toolkit, 'cuda') && runner.arch != 'arm64' }}
shell: bash
run: sudo apt-get install -y libnccl2 libnccl-dev
- name: CUDA packages and driver report - name: CUDA packages and driver report
if: ${{ startsWith(inputs.toolkit, 'cuda') }} if: ${{ startsWith(inputs.toolkit, 'cuda') }}