NVIDIA GPU 36 — DGX Spark Development Workflow

00

Topics We'll Cover

Where Spark Fits in Your Workflow
Connecting from Your Laptop — VS Code Remote SSH
Containers as the Default Environment
Python on ARM64 — Wheels Reality
PyTorch & Friends on Spark
JupyterHub for Multi-User Spark
Fine-Tuning — QLoRA, FSDP, NeMo
Dataset Storage & Streaming
Quantisation Pipeline — HF → NVFP4 / MX-FP4 NIM
Custom Kernels with Triton on Spark
Spark Local + Cloud Burst Pattern
Interactive: Pick Your Workflow

01

Where Spark Fits in Your Workflow

Spark is best as the middle tier of your AI development stack:

Laptop (the IDE)

Your Mac or Linux laptop is just an editor + browser. Don't run heavy local Python kernels — SSH into Spark for everything.

Spark (the workstation)

Where models run. Where you fine-tune. Where you iterate. The ~3-second feedback loop on a 7 B inference test makes prompt-engineering practical.

Cloud (the burst tier)

For pretraining 7 B+, full BF16 ablations, > 70 B-class fine-tunes, big multi-GPU jobs. Rent H100s for hours, not months.

The pattern: iterate on Spark, scale on cloud. Your code lives in git, runs in a container that's identical on Spark and on a cloud DGX H100. The only code change between local and cloud is the model checkpoint path.

02

Connecting from Your Laptop — VS Code Remote SSH

VS Code Remote SSH on Spark is excellent — the X925 cores are fast enough that the language server runs in real time, and ARM64 is fully supported by the VS Code server since 1.75.

~/.ssh/config on your laptop

Host spark
    HostName spark.local         # or 100.64.x.y if Tailscale
    User brendan
    ForwardAgent yes
    ServerAliveInterval 30
    LocalForward 8000 localhost:8000     # vLLM endpoint forwarded
    LocalForward 8888 localhost:8888     # Jupyter forwarded
    LocalForward 3000 localhost:3000     # Open WebUI forwarded

Then in VS Code: Remote-SSH: Connect to Host... → spark. Open ~/code in the workspace. Install Python, Pylance, GitLens, and Dev Containers extensions on the remote (one-time). The first connection takes ~90 s to install the VS Code server; subsequent connects are instant.

If you're outside your home LAN, add Tailscale: Spark gets a stable hostname like spark.tail-scale.ts.net, accessible from anywhere with WireGuard.

03

Containers as the Default Environment

Every NVIDIA-blessed dev workflow on Spark is container-based. Two complementary patterns:

One-shot container

Disposable: spin up, run a script, throw away.

docker run --rm -it --gpus all \
    -v $PWD:/work -w /work \
    nvcr.io/nvidia/pytorch:24.10-py3 \
    python train.py

Long-lived dev container

VS Code Dev Containers: same as one-shot but kept alive while you edit.

// .devcontainer/devcontainer.json
{
  "image": "nvcr.io/nvidia/pytorch:24.10-py3",
  "runArgs": ["--gpus", "all"],
  "mounts": ["source=/opt/models,target=/models,type=bind"]
}

Pin to ARM64 manifests

Always check that the image you're pulling has an ARM64 variant: docker manifest inspect nvcr.io/... should list linux/arm64. The good news: every NGC image since 2024 ships ARM64 alongside x86. Third-party Docker Hub images often don't — budget time to switch or rebuild.

04

Python on ARM64 — Wheels Reality

The pre-2023 ARM64 wheel-availability anxiety is mostly resolved in 2026. Major scientific Python packages ship aarch64 wheels: numpy, scipy, pandas, torch, jax, tensorflow, transformers, accelerate, peft, datasets, sentence-transformers, vLLM, etc.

uv is the recommended environment manager: faster than poetry, native ARM64 binary, handles CUDA index automatically.
Conda / Mamba works fine; conda-forge has full ARM64 coverage.
Pip wheel cache: ~/.cache/pip on Spark; warm it once with the heavy deps, all later venvs are fast.
The exceptions: niche older packages (old-package-1.x) without ARM64 wheels — pip install will try to compile, often fails. Either upgrade to a newer release that has wheels or live with the build.

canonical Spark Python env

curl -fsSL https://astral.sh/uv/install.sh | sh
uv venv -p 3.12 .venv
source .venv/bin/activate
uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu130
uv pip install transformers accelerate peft datasets bitsandbytes vllm

05

PyTorch & Friends on Spark

PyTorch on Spark ARM64 + CUDA 13.0 is supported via the cu130 aarch64 wheels (GB10 is sm_121, which CUDA 12.x builds do not target). Compile-once-run-anywhere: code that works on Spark works on a cloud H100 with the same wheel index.

torch.compile works; first compile takes longer on Spark than on a beefy x86 + 5090 (CPU side is slower), but caches normally.
Flash-Attention: flash-attn 2.6+ has ARM64 wheels. Flash-Attn 3 is Hopper-only (H100/H800); Blackwell support is in FlashAttention-4.
bitsandbytes: ARM64 wheel from 0.43; 8-bit optimisers and 4-bit quantised models work.
Triton: bundled with PyTorch; native ARM64 from Triton 3.0. Custom kernels portable.
JAX: jax[cuda12] wheel works; less common in this stack but supported.
TensorFlow: NGC image (nvcr.io/nvidia/tensorflow) has ARM64 build; standalone pip install tensorflow historically lagged.

torch.compile on Spark

The Inductor backend compiles to Triton kernels which run on Spark's Blackwell GPU at full speed. Real-world: torch.compile on a Llama fine-tune gives ~15–25% speedup on Spark, similar to a discrete RTX 5090.

06

JupyterHub for Multi-User Spark

One Spark for a small team works well with JupyterHub. Each user gets their own kernel; the unified-memory architecture means model loading is fast and shared.

JupyterHub on Spark with The Littlest JupyterHub

curl -L https://tljh.jupyter.org/bootstrap.py | \
    sudo -E python3 - --admin brendan

# add NVIDIA-flavoured kernel image
sudo tljh-config set user_environment.default_app jupyterlab
sudo tljh-config set images.image nvcr.io/nvidia/pytorch:24.10-py3
sudo tljh-config reload

Memory & GPU sharing patterns:

Default: all users share one CUDA context; one user's OOM kills the GPU for everyone.
Better: launch each user's kernel in its own --gpus all container; CUDA MPS daemon coordinates.
Best for production: don't host JupyterHub on a Spark with a serving workload — let Jupyter users have one Spark, serving has another.

07

Fine-Tuning — QLoRA, FSDP, NeMo

What you can fine-tune on a single Spark depends on memory, not compute. Rough budgets:

Approach	Model size feasible	Tool	Notes
Full BF16 fine-tune	up to ~7 B	HF Trainer / NeMo SFT	weights+grad+Adam state ~16 bytes per param (mixed precision): 7 B → ~112 GB
BF16 + FSDP CPU offload	up to ~13 B	HF + Accelerate	slow; CPU-GPU traffic over unified-mem fabric
LoRA on BF16 base	up to ~32 B	peft + transformers	frozen weights at BF16; only adapters train
QLoRA (NF4 base + BF16 LoRA)	up to ~70 B	peft + bitsandbytes	the canonical Spark fine-tuning recipe
NeMo SFT	up to ~70 B with DPO	NeMo + NeMo Aligner	preferred for reproducible large-model alignment
QLoRA on two-Spark pair	up to ~120 B	peft + Accelerate FSDP	experimental; PP=2 with Accelerate works

QLoRA Llama-3.3-70B on a single Spark

accelerate launch \
    --num_processes 1 --mixed_precision bf16 \
    train.py \
    --model_name meta-llama/Llama-3.3-70B \
    --load_in_4bit --bnb_4bit_quant_type nf4 \
    --lora_r 16 --lora_alpha 32 \
    --gradient_checkpointing \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --num_train_epochs 1

08

Dataset Storage & Streaming

4 TB NVMe SKU recommended for serious dataset work. Models pile up: 35 GB per 70 B-class checkpoint, × several variants, plus datasets.
HF datasets: load_dataset(streaming=True) works around capacity for large corpora. Apache Arrow over the shared LPDDR5x is fast.
NFS mount: the 10 GbE RJ-45 port limits NFS reads to ~1 GB/s; OK for episodic dataset access, slow for steady training.
200 GbE NAS: a TrueNAS box on the QSFP link delivers ~6 GB/s — useful if you have one.
Cloud-side data lake: S3 or Lustre via boto3 over Wi-Fi 7; ~300–500 MB/s real-world. Stream training samples on demand.

Pre-shard your data

The Spark IO pipeline benefits massively from pre-tokenised, pre-shuffled, pre-packed Parquet shards on local NVMe. datasets.save_to_disk() + sharding by 1 GB pieces makes the dataloader's life much easier than streaming raw text.

09

Quantisation Pipeline — HF → NVFP4 / MX-FP4 NIM

You fine-tuned in BF16 / QLoRA. Now you want to deploy in FP4 to a 7 tok/s decode throughput. Blackwell's 5th-gen tensor cores accelerate both NVFP4 and MX-FP4; NVFP4 (16-element blocks, E4M3 scale + per-tensor FP32) is the default in TensorRT-LLM and Transformer Engine, while MX-FP4 (32-element blocks, E8M0 scale) follows the OCP MX spec. The pipeline:

Merge LoRA back into the base weights: peft.merge_and_unload() → full BF16 model on disk.
Convert to FP4 with TensorRT-LLM: quantize.py --kv_cache_dtype fp8 --qformat nvfp4 (or mx_fp4) — calibrates per-microblock scales on a small calibration set.
Build TRT-LLM engine: trtllm-build --checkpoint_dir ./calib --quantization nvfp4 --output_dir ./engine — ~15–25 minutes on Spark for a 70 B.
Wrap as NIM: NVIDIA's nim-builder tool packages a TRT-LLM engine into an OpenAI-compatible NIM container.
Run: docker run --gpus all -p 8000:8000 my-org/llama3.3-70b-mxfp4-nim:v1.

Once you've done this once, you have a reproducible private NIM that ships your fine-tuned model in production-grade form. Same artifact runs on a cloud H100 if you ever burst.

10

Custom Kernels with Triton on Spark

Triton works fully on Spark's Blackwell GPU. The compiler emits PTX targeting GB10's sm_121; you get tensor-core acceleration with a Python-like kernel DSL.

Triton matmul on Spark

import triton, triton.language as tl

@triton.jit
def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid = tl.program_id(0)
    # ... tile-based matmul body ...
    tl.store(c_ptr + offsets, accumulator)

Useful patterns Spark-side:

FlashAttention variants — rewrite-and-tune for your model's exact head dim / context shape.
Custom fused activations — SwiGLU + RMSNorm in one kernel for a tight fine-tune loop.
Sparse attention for very long context.
Per-microblock dequant kernels if your custom MX-FP4 layout differs from the standard one.

Profile with nsys profile — same workflow as on x86. Triton's autotuner takes longer on Spark's CPU but produces equivalent kernels.

11

Spark Local + Cloud Burst Pattern

The most productive workflow: do all your iteration on Spark, burst to the cloud only for the heavy run.

Iterate on Spark

Code in your ~/code/project on Spark.
QLoRA fine-tune the 7–13 B variant of your target model with your tiny dataset to validate the script end-to-end.
Inspect outputs in JupyterLab; tweak; commit.
Total cost: power bill + your time.

Burst to cloud

Push branch, trigger CI to run the same script in a cloud DGX H100 (e.g. Lambda Labs, RunPod, OCI).
Same container image, same code; just bigger model, full dataset, more epochs.
Spend 10–100 GPU-hours, get full-precision result back.
Total cost: ~$30–$300 per run.

Why this beats "just always use cloud": iteration cycles. A 30 s sanity-check run on Spark is essentially free; the same on rented cloud GPUs costs $1 in spin-up time alone. You're more willing to experiment when each experiment is cheap.

12

Interactive: Pick Your Workflow

Task

Team size