# DGX Spark FlashAttention wheel

This repository hosts a DGX Spark compatible FlashAttention wheel used by the PII fine-tuning tutorial stack.

The wheel was taken from the known-good `unsloth-bench` environment on DGX Spark Alice and is intended for:

- NVIDIA DGX Spark / Linux `aarch64`
- Python 3.12 / CPython `cp312`
- CUDA 13 PyTorch stack
- `flash_attn==2.8.3.post1`

Wheel URL:

```bash
https://huggingface.co/N8Programs/dgx-flash-attn-wheel/resolve/main/flash_attn-2.8.3.post1-cp312-cp312-linux_aarch64.whl
```

## 1. Create a clean conda environment

On the DGX Spark host:

```bash
conda create -n piiEnv python=3.12 pip -y
conda activate piiEnv
python -m pip install --upgrade pip
```

Use `python -m pip` from inside the conda environment. If `pip install` reports an externally-managed-environment error, the environment is probably missing its own Python/pip and is falling back to system Python.

## 2. Install the CUDA 13 PyTorch base

```bash
python -m pip install \
  --extra-index-url https://download.pytorch.org/whl/cu130 \
  torch==2.12.0+cu130 \
  triton==3.7.0 \
  packaging==26.2 \
  ninja==1.13.0 \
  psutil==7.2.2
```

## 3. Install the DGX Spark FlashAttention wheel

```bash
python -m pip install \
  https://huggingface.co/N8Programs/dgx-flash-attn-wheel/resolve/main/flash_attn-2.8.3.post1-cp312-cp312-linux_aarch64.whl
```

This avoids rebuilding FlashAttention from source. On DGX Spark, source builds can be slow and easy to misconfigure because FlashAttention may try to compile multiple CUDA architectures unless carefully constrained.

## 4. Install the training stack

```bash
python -m pip install \
  transformers==5.5.0 \
  accelerate==1.13.0 \
  peft==0.19.1 \
  trl==0.24.0 \
  datasets==5.0.0 \
  bitsandbytes==0.49.2 \
  cut-cross-entropy==25.1.1 \
  sentencepiece==0.2.1 \
  tokenizers==0.22.2 \
  safetensors==0.8.0 \
  protobuf==7.35.0 \
  numpy==2.4.6 \
  pandas==3.0.3 \
  Pillow==12.2.0 \
  PyYAML==6.0.3 \
  hf_transfer==0.1.9 \
  msgspec==0.21.1 \
  tqdm==4.68.2 \
  tyro==1.0.13 \
  wandb==0.27.2
```

## 5. Install Unsloth without dependency resolution

The known-good DGX Spark stack uses `torch==2.12.0+cu130`, but the current Unsloth package metadata may declare an older Torch upper bound. Preserve the working CUDA stack by installing Unsloth without dependencies:

```bash
python -m pip install --no-deps unsloth_zoo==2026.6.2 unsloth==2026.6.2
```

Pip may warn about dependency metadata mismatches if you later inspect the environment. The important part is that Torch, FlashAttention, Transformers, Unsloth, and CUDA import together successfully.

## 6. Smoke test

```bash
python - <<PY
import torch
import transformers
import flash_attn
import unsloth
import unsloth_zoo

print("torch", torch.__version__, "cuda", torch.version.cuda, "available", torch.cuda.is_available())
print("transformers", transformers.__version__)
print("flash_attn", flash_attn.__version__)
print("unsloth", unsloth.__version__)
print("unsloth_zoo", unsloth_zoo.__version__)
PY
```

Expected shape of a healthy DGX Spark setup:

```text
torch 2.12.0+cu130 cuda 13.0 available True
transformers 5.5.0
flash_attn 2.8.3.post1
unsloth 2026.6.2
unsloth_zoo 2026.6.2
```

## Optional one-shot installer

```bash
cat > install_dgx_spark_env.sh <<SH
#!/usr/bin/env bash
set -euo pipefail

python -m pip install --upgrade pip
python -m pip install \
  --extra-index-url https://download.pytorch.org/whl/cu130 \
  torch==2.12.0+cu130 triton==3.7.0 packaging==26.2 ninja==1.13.0 psutil==7.2.2
python -m pip install \
  https://huggingface.co/N8Programs/dgx-flash-attn-wheel/resolve/main/flash_attn-2.8.3.post1-cp312-cp312-linux_aarch64.whl
python -m pip install \
  transformers==5.5.0 accelerate==1.13.0 peft==0.19.1 trl==0.24.0 datasets==5.0.0 \
  bitsandbytes==0.49.2 cut-cross-entropy==25.1.1 sentencepiece==0.2.1 tokenizers==0.22.2 \
  safetensors==0.8.0 protobuf==7.35.0 numpy==2.4.6 pandas==3.0.3 Pillow==12.2.0 \
  PyYAML==6.0.3 hf_transfer==0.1.9 msgspec==0.21.1 tqdm==4.68.2 tyro==1.0.13 wandb==0.27.2
python -m pip install --no-deps unsloth_zoo==2026.6.2 unsloth==2026.6.2
SH
bash install_dgx_spark_env.sh
```