Skip to content

GPU.md

Latest commit

 

History

History
35 lines (25 loc) · 2.44 KB

GPU.md

File metadata and controls

35 lines (25 loc) · 2.44 KB
Raw

GPU deployment

Nvidia GPU support is available for analysis task in the worker process. This can significantly speed up processing of tracks.

ARM (DGX Spark / GB10) support: the -nvidia-arm image adds support for NVIDIA GPUs on ARM64 hosts, such as the DGX Spark and other GB10-based machines, for both analysis and clustering. This image is EXPERIMENTAL. Use it the same way as the regular -nvidia image, just pick the -nvidia-arm tag.

We suggest 8GB VRAM on GPU, with less you can experience the NON BLOCKING OutOFMemory error (that are handled by switching to CPU). The PER_SONG_MODEL_RELOAD env variable, that by default is TRUE, help cleaning the memory by entirely reloading the model each time, on the other side it slow the analysis process.

GPU-accelerated clustering is also available through RAPIDS cuML. It can give a 10-30x speedup on clustering tasks.

Features:

  • GPU-accelerated KMeans, DBSCAN, PCA and SpectralClustering using RAPIDS cuML
  • Automatic fallback to CPU if the GPU is unavailable or hits an error
  • Works with all existing clustering configurations and parameters
  • Compatible with NVIDIA GPUs on CUDA 13 or later (*)

(*) CUDA 13 raises the minimum host driver to >=580.x (up from >=570.x for the previous CUDA 12.8 image) - older drivers are NOT supported by the published build, but you can try to build your own image as described in NeptuneHub#265

To enable GPU clustering:

  1. Use the NVIDIA image (for example nvidia/cuda:13.3.1-cudnn-runtime-ubuntu24.04)
  2. Set the value in your .env file, or in the Setup Wizard: 
    USE_GPU_CLUSTERING=true
  3. Make sure the NVIDIA Container Toolkit is installed on the host
  4. Use the GPU compose file deployment/docker-compose-nvidia.yaml. A worker-only GPU example is kept in deployment/test/docker-compose-nvidia-worker-test.yaml

Notes:

  • GMM stays on CPU, there is no cuML implementation for it
  • Spectral clustering runs on cuML only with assign_labels='kmeans' and a nearest_neighbors / precomputed affinity (what the clustering search uses); any other combination falls back to scikit-learn
  • GPU clustering is disabled by default (USE_GPU_CLUSTERING=false)
  • The GPU is also used by the audio analysis models (ONNX inference)
  • The index build and the similarity queries are not GPU accelerated; they are IO bound rather than compute bound, see ALGORITHM