Dockerfile

# Baseline image for DGX Spark (GB10, sm_121a) — stock Unsloth recipe.
#
# This is the *unoptimized baseline*. The whole point of this project is to
# measure speedups *on top of* this image, so do not add Liger / FA3 / TE /
# torch.compile here. Optimized variants go in docker/Dockerfile.opt* later.
#
# Pattern follows the NVIDIA DGX Spark playbook + haven-jeon/unsloth-vllm-gb10:
#   - CUDA 13.0 (via NGC pytorch:25.10-py3)
#   - Triton built from source at commit with sm_121a support
#   - xformers built from source with TORCH_CUDA_ARCH_LIST=12.1
#   - Unsloth + unsloth_zoo at latest at build time, then pip-frozen for repro

# NGC 26.03 is what this project standardised on after a smoke test sweep.
# Tried 25.10 first (per haven-jeon recipe), but its bundled inductor is
# incompatible with Triton 3.6.0's KernelMetadata: torch.compile crashes at
# step 0 with "'KernelMetadata' object has no attribute 'cluster_dims'".
# 26.03's inductor handles it correctly.
FROM nvcr.io/nvidia/pytorch:26.03-py3

ENV CUDA_HOME=/usr/local/cuda \
    CUDA_PATH=/usr/local/cuda \
    PATH=/usr/local/cuda/bin:$PATH \
    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
    C_INCLUDE_PATH=/usr/local/cuda/include \
    CPLUS_INCLUDE_PATH=/usr/local/cuda/include \
    TORCH_CUDA_ARCH_LIST="12.1" \
    TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas \
    UV_SYSTEM_PYTHON=1 \
    PIP_ROOT_USER_ACTION=ignore

RUN python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}'); assert torch.version.cuda is not None"

# ── Triton from source @ v3.6.0 (sm_121a support) ──
# Earlier attempts pinned to commit c5d671f from the haven-jeon recipe; the
# v3.6.0 release tag is in the same lineage and is what our existing
# unsloth-dgx-spark image successfully uses on this hardware.
RUN git clone https://github.com/triton-lang/triton.git /tmp/triton && \
    cd /tmp/triton && \
    git checkout v3.6.0 && \
    pip install -r python/requirements.txt && \
    pip install . && \
    rm -rf /tmp/triton

# ── xformers v0.0.33 from source (sm_121a) ──────────────────────────────────
RUN MAX_JOBS=8 TORCH_CUDA_ARCH_LIST=12.1 \
    pip install --no-deps --no-build-isolation \
        "git+https://github.com/facebookresearch/xformers@v0.0.33"

# ── torchao conflicts with this stack (per haven-jeon notes) ────────────────
RUN pip uninstall -y torchao || true

# ── Unsloth ecosystem, baseline pins ────────────────────────────────────────
# bitsandbytes: 0.49+ is the first line with aarch64 + CUDA 13 wheels on PyPI
# (older haven-jeon recipes pinned 0.45.5 — no longer available for this arch).
RUN pip install --no-deps "bitsandbytes>=0.49.0"
RUN pip install --no-deps unsloth unsloth_zoo
# transformers: Unsloth 2026.5+ refuses Qwen3.5 below 5.2.0 (ImportError on FastModel.from_pretrained).
# haven-jeon's original 4.57.6 pin is fine for Llama 3.2 but breaks Qwen3.5,
# which is the secondary-case-study model. 5.5.x is the proven Qwen3.5 line
# for this stack.
RUN pip install \
        "transformers==5.5.0" \
        "trl==0.24.0" \
        "peft>=0.18.0" \
        "datasets>=3.4.1,<4.4.0" \
        "accelerate" \
        "safetensors" \
        "huggingface_hub" \
        "hf_transfer" \
        "tyro" \
        "sentencepiece"

# ── Freeze versions for reproducibility ─────────────────────────────────────
# We don't import unsloth / bitsandbytes here — unsloth_zoo refuses to import
# without a torch accelerator visible, and `docker build` runs without GPU.
# Those imports are exercised in scripts/check_image.sh (run with --gpus=all).
RUN pip freeze > /opt/pip-freeze-baseline.txt && \
    python -c "import torch, triton, transformers, trl, peft; \
        print(f'torch={torch.__version__}'); \
        print(f'triton={triton.__version__}'); \
        print(f'transformers={transformers.__version__}'); \
        print(f'trl={trl.__version__}'); \
        print(f'peft={peft.__version__}')"

WORKDIR /workspace
CMD ["/bin/bash"]