DGX Spark/Guides/PyTorch on DGX Spark
Guide · November 2026

PyTorch on a DGX Spark: CUDA 13 aarch64 wheels, the sm_121 warning, NGC containers and compiled packages

By Samuel Seidel · Published November 7, 2026 · Sources checked November 7, 2026

Stock PyTorch runs on a DGX Spark: since PyTorch 2.11, pip install torch on Linux aarch64 pulls a CUDA wheel, and as of 7 November 2026 the newest is torch 2.15.0 (28 October) with CUDA 13.2, while CUDA 13.0 builds, which match DGX OS's CUDA 13.0.3 and driver 580.178.04, end at torch 2.14.1. The wheels carry sm_120 kernels that run natively on the GB10's sm_121, so the old "capability 12.1 ... (8.0) - (12.0)" warning is a false alarm. The weak spots are compiled extras: bitsandbytes 0.50.2 and triton 3.9.0 have aarch64 wheels, while flash-attn and xformers still need source builds or NVIDIA's NGC container.

PackageLatest (date)aarch64 wheelGB10 (sm_121) statusSource
torch2.15.0 (2026-10-28)Yes: cu132 on PyPI, cu134 index; cu130 up to 2.14.1Works; sm_120 SASS runs on sm_1212.15 notes
triton3.9.0 (2026-10-27)YesWorks without env vars from 3.6 on (reported)issue 10331
bitsandbytes0.50.2 (2026-08-27)YesCUDA 13.0 to 13.2 aarch64 build includes sm121bnb docs
torchaudio2.11.1 (2026-10-28)YesWorks; maintenance releases onlyreleases
flash-attn2.8.3.post2 (2026-10-14)No (source only)sm121 issue open since 2025-10-28issue 1969
flash-attn-44.0.0b38 (2026-11-04)Pure Pythonsm_120 integration PR still openPR 2634
xformers0.0.35 (2026-02-20)No (x86_64 and Windows only)Source buildPyPI, README

Versions and wheel platforms read from PyPI and download.pytorch.org on 7 November 2026; status from the linked first-party pages and issues.

Which PyTorch wheel should you install on a DGX Spark?

Either torch 2.14.1 from the cu130 index, which matches the CUDA 13.0 the Spark ships, or torch 2.15.0 from PyPI if you want the current release. The 2.11 notes made aarch64 CUDA wheels the PyPI default; before that, a plain pip install on an Arm box could leave you on a CPU-only build. The 2.15 notes then moved the default from CUDA 13.0 to 13.2, publish cu132 and cu134 wheels, stopped building cu126 and cu130, and require Python 3.11 or newer. The cu130 index has 2.14.1 aarch64 wheels for Python 3.10 to 3.15, including the free-threaded builds.

pip install uv
uv venv --python 3.12 .venv && source .venv/bin/activate

# CUDA 13.0, same minor version as DGX OS 7.6.0
uv pip install torch==2.14.1 torchvision --index-url https://download.pytorch.org/whl/cu130

# or the current release, CUDA 13.2 wheel from PyPI
# uv pip install torch==2.15.0 torchvision

python - <<'EOF'
import torch
print(torch.__version__, torch.version.cuda)
print(torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
print(torch.cuda.get_arch_list())
x = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
print((x @ x).float().abs().mean().item())
EOF

The last two lines should print an arch list ending in sm_120 and a number, not an error. Both builds work on the Spark's driver: NVIDIA's minor version compatibility page sets 580 as the minimum for all of CUDA 13.x, and the DGX Spark release notes list driver 580.178.04. The same page has the catch: "Applications that compile device code to PTX will not work on older drivers." Anything that JIT-compiles PTX from a 13.2 or 13.4 toolkit on a 13.0 driver can fail where cu130 would not. That is why I default to 2.14.1+cu130 on a Spark and move to 2.15 when a package needs it.

On a rented GPUwerk Spark the same commands work as root inside the container, where the driver comes from the host (see the tenant environment). Create the venv under /workspace, because only that directory survives a rebuild; JupyterLab installs into it the same way.

What does the sm_121 warning mean?

Nothing is wrong. Up to PyTorch 2.10, importing torch on a GB10 printed:

UserWarning:
    Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.
    Minimum and Maximum cuda capability supported by this version of PyTorch is
    (8.0) - (12.0)

PyTorch maintainer ptrblck answered on the PyTorch forum in October 2025: "Ignore the warning as it's wrong since sm_121 is SASS binary compatible with sm_120." In issue 172629 a GB10 owner showed in September 2026 that torch 2.10.0+cu130 still printed it, with an arch list of sm_80, sm_90, sm_100, sm_110, sm_120 and compute_120, while 2.13.0+cu130 printed nothing because its check accepts 12.1 within the 12.x range. The 2.15 notes add one sm_121 item: "Reuse the CUTLASS sm120 kernel generators on sm121 GPUs."

Whether a native sm_121a build would be faster is an open claim. PR 188191, submitted by an AI agent on a user's behalf, argued for adding 12.1a to the default arch list, citing a "25-30% throughput penalty" and community nanochat MFU of 5.6 to 5.9 against 4.3 to 4.7. It was closed without merging on 25 September 2026. I have found no maintainer confirmation of that gap. If you want to test it, the route is a source build with TORCH_CUDA_ARCH_LIST="12.1a", which can take hours on the Spark's 20 cores.

NGC PyTorch container or pip?

On a Spark you own, start with the container; it is the build NVIDIA tests on GB10 and it bundles the extras that are hard to compile. NVIDIA's 26.10 release notes list torch 2.15.0a0+875d815502, CUDA 13.4.3, Transformer Engine 2.20, TensorRT 11.4, JupyterLab 4.6.4, Python 3.12 on Ubuntu 24.04, and from this release a FlashAttention 4 beta (4.0.0b24) next to the existing FlashAttention 2. The arm64 image is a 12.7 GB download by my sum of its nvcr.io manifest layers.

docker pull nvcr.io/nvidia/pytorch:26.10-py3
docker run --gpus all -it --rm --ipc=host \
  -v "$HOME/.cache/huggingface:/root/.cache/huggingface" \
  -v "${PWD}:/workspace" -w /workspace \
  nvcr.io/nvidia/pytorch:26.10-py3

That is the run line from NVIDIA's PyTorch fine-tuning playbook, which itself still pins 25.11; I swapped in the current tag. Three things to know. The container ships a pip constraints file at /etc/pip/constraint.txt that stops pip from replacing its torch; edit it if you must change a pinned package. NVIDIA stopped producing separate -igpu images with 26.03, so use the plain tag. And the NGC page in NVIDIA's own Spark user guide still shows pytorch:24.08-py3, a CUDA 12.6 image from before NVIDIA's containers were "optimized for Blackwell" (from 25.01, per the release notes); do not copy that tag.

Is the container faster? The only number I found is old: in the NVIDIA forum thread "Effective PyTorch and CUDA", an owner timed a ViT model at over 7 minutes with the cu130 pip wheel against under 4 minutes in the NGC container, at about 96 percent GPU use either way (October 2025, PyTorch 2.9 era). Treat that as a reason to compare on your own job, not as a current ratio. Docker setup on DGX OS is in the sibling guide to Docker and NGC containers; a rented instance has no Docker socket, so there pip is the only option.

Which compiled packages work on aarch64 and sm_121?

Triton, bitsandbytes and torchaudio install from wheels; flash-attn and xformers do not.

Triton

Triton 3.9.0 (27 October 2026) has manylinux aarch64 wheels, and the PyTorch 2.15 notes refer to Triton 3.9. Much of the advice online is out of date. In Triton issue 10331 a contributor showed in August 2026 that Triton has shipped a second "ptxas-blackwell" binary since 3.6.0 that assembles sm_121a, verified kernels end to end on 3.7.1 with no environment variables, and pointed out that the widely copied TRITON_PTXAS_PATH has no effect on Blackwell targets; the override for that binary is TRITON_PTXAS_BLACKWELL_PATH. The same thread warns that TRITON_OVERRIDE_ARCH=sm90 builds kernels the driver rejects, or worse, that run and compute wrong numbers. To check which ptxas your install uses:

python -c "from triton.backends.nvidia.compiler import get_ptxas; print(get_ptxas(121).path, get_ptxas(121).version)"

flash-attn and FlashAttention 4

flash-attn 2.8.3.post2 publishes no wheels on PyPI, and issue 1969 ("Support for compute capability 12.1") has been open since 28 October 2025. Users in it report working source builds; the commands posted there are:

git clone --depth=1 https://github.com/Dao-AILab/flash-attention ./flash-attention
cd ./flash-attention
export MAX_JOBS=4 NVCC_THREADS=2        # the thread used 16; one user needed 4
FLASH_ATTENTION_FORCE_BUILD="TRUE" \
FLASH_ATTENTION_FORCE_CXX11_ABI="FALSE" \
FLASH_ATTENTION_SKIP_CUDA_BUILD="FALSE" \
pip3 wheel . -v --no-deps -w ./wheels/
pip3 install ./wheels/flash_attn*.whl

Two shortcuts. In April 2026 a user reported that FlashAttention 3 built in the NGC 26.03 container in about 50 minutes and then failed with "no kernel image is available" because it compiled only sm80 and sm90 kernels, while the FlashAttention 2.7.4 already in that container ran on GB10 with outputs bit-exact to PyTorch's scaled_dot_product_attention. And SDPA alone is enough for many models: several commenters fell back to it. FlashAttention 4 (flash-attn-4 4.0.0b38 on PyPI) is CuTe DSL code, but its consumer-Blackwell port, PR 2634, was still open on 7 November 2026.

xformers, bitsandbytes, torchaudio

xformers 0.0.35 ships no aarch64 wheel, so build it as its README describes, with the arch list set for GB10 and ninja parallelism capped to avoid running out of memory:

pip install ninja
export TORCH_CUDA_ARCH_LIST="12.1" MAX_JOBS=4
pip install -v --no-build-isolation -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

bitsandbytes is the good news: its installation guide lists Linux aarch64 wheels for CUDA 13.0 to 13.2 compiled for sm75 through sm121, so pip install bitsandbytes is all QLoRA needs. (Its separate note telling "NVIDIA Spark devices (such as an RTX Spark laptop)" to take PyTorch from pypi.nvidia.com reads, to me, as aimed at the Windows-on-Arm laptops, since Linux aarch64 CUDA wheels are on PyPI.) torchaudio 2.11.1 has aarch64 wheels and works with torch 2.11 or newer; its maintainers say releases since 2.10 add no features after the migration tracked in issue 3902, so plan new audio code around something else.

How do you stop a training job from freezing the Spark?

Cap PyTorch's allocator, because on GB10 GPU memory is host memory and an out-of-memory error can arrive late or not at all. PyTorch issue 201023 (open, filed 15 October 2026 on torch 2.14.0+cu130) lays out the mechanism: cudaMalloc on GB10 does not fail until Linux has evicted the page cache and swapped out processes such as sshd, and the caching allocator's garbage_collection_threshold is silently ignored unless a per-process memory fraction below 1.0 is also set. The earlier full-system hang in issue 174358 was closed in April 2026 after PyTorch's team could no longer reproduce it on kernel 6.17 with PyTorch 2.11. Based on issue 201023, I set both:

export PYTORCH_ALLOC_CONF=expandable_segments:True,garbage_collection_threshold:0.8

# at the top of the training script
import torch
torch.cuda.set_per_process_memory_fraction(0.85)   # also makes the GC threshold take effect

By my arithmetic, 0.85 of the roughly 121.7 GB a Spark exposes is about 103 GB for PyTorch, leaving some 18 GB for the OS, data loaders and an SSH session. NVIDIA's July 2026 software release also improved out-of-memory handling and added user feedback under memory pressure, per the DGX Spark release notes. For what fits at all, the memory calculator works for training footprints too.

How fast does PyTorch train on a Spark?

NVIDIA's own launch numbers are the only systematic set, and they are peaks. Its October 2025 performance post reports, as (batch size x steps x sequence length) divided by training time:

ModelMethodPeak tokens/sSourceDate
Llama 3.2 3BFull fine-tuning82,739.2NVIDIA blog2025-10-24
Llama 3.1 8BLoRA53,657.6NVIDIA blog2025-10-24
Llama 3.3 70BQLoRA5,079.4NVIDIA blog2025-10-24

NVIDIA's figures. GPUwerk has published no training benchmarks.

The fine-tuning playbook has scripts for the same three cases (Llama3_3B_full_finetuning.py, Llama3_8B_LoRA_finetuning.py, Llama3_70B_qLoRA_finetuning.py) and installs transformers 4.57.1, peft 0.18.0, trl 0.25.1 and bitsandbytes into the container. Time per epoch on your data is easy to estimate from those rates: a 50-million-token LoRA epoch on the 8B is about 16 minutes at 53,658 tokens/s, and the same tokens through a 70B QLoRA take about 2.7 hours at 5,079 tokens/s, by my division. What to fine-tune and how to judge the result is in fine-tuning on your own data; the 32 GB alternative is in fine-tuning on an RTX 5090.

My recommendation: on a Spark you own, develop in nvcr.io/nvidia/pytorch:26.10-py3 and pip only what is missing; on a rented one, torch 2.14.1+cu130 in a /workspace venv, bitsandbytes from PyPI, SDPA instead of flash-attn until issue 1969 closes. For serving rather than training, the vLLM and SGLang guides install their own pinned torch, so keep them in separate environments.

FAQ

Does pip install torch work on a DGX Spark?

Yes. Since PyTorch 2.11, pip install torch on Linux aarch64 installs a CUDA wheel from PyPI. PyTorch 2.15.0 (28 October 2026) ships CUDA 13.2 wheels by default and CUDA 13.4 wheels on its own index. CUDA 13.0 wheels, matching DGX OS's CUDA 13.0.3, stop at 2.14.1, installed from download.pytorch.org/whl/cu130.

Why does PyTorch say my GB10 with capability 12.1 is not supported?

That warning is wrong. PyTorch 2.10 and earlier printed that the supported range was 8.0 to 12.0, but PyTorch maintainer ptrblck explained that sm_121 is SASS binary compatible with sm_120, which the wheels include. A GB10 owner reported in September 2026 that torch 2.13.0+cu130 no longer prints it. The kernels run natively.

Does flash-attention work on a DGX Spark?

Not from pip. flash-attn 2.8.3.post2 has no prebuilt wheels on PyPI, and the GitHub issue asking for sm_121 support has been open since 28 October 2025. Users report source builds that work, NVIDIA's NGC PyTorch 26.10 container ships FlashAttention 2 plus a FlashAttention 4 beta, and PyTorch's built-in scaled_dot_product_attention works without it.

Should I use the NGC PyTorch container or pip on a Spark?

For a Spark you own, the container. NGC PyTorch 26.10 bundles torch 2.15.0a0, CUDA 13.4.3, Transformer Engine 2.20 and FlashAttention, and in October 2025 one owner timed a ViT job at under 4 minutes in the container against over 7 minutes with pip wheels. On a rented GPUwerk Spark there is no Docker, so pip it is.

Why does my DGX Spark freeze instead of raising CUDA out of memory?

Because GPU memory is system memory. On GB10, cudaMalloc takes host RAM and does not fail until the kernel has evicted page cache and swapped processes, per PyTorch issue 201023, which is still open. PyTorch maintainers could no longer reproduce the full hang on the April 2026 stack; capping the allocator with set_per_process_memory_fraction is the cheap protection.

Sources

Every figure above is attributed to its source in the sentence or table where it appears. Pages were read on 7 November 2026 unless a date says otherwise.

First top-up: pay $20, get $30 in credit

Try your training script on a GB10 first.

A rented 128 GB Spark takes the same aarch64 pip wheels as one you own, so you can find out which of your compiled dependencies build before you buy hardware.

A dedicated DGX Spark in the EU, $0.79/hour, billed per minute. $30 of credit is about 37 hours of compute.

Create account Fine-tuning on a Spark