unsloth/docker/setup_qemu.sh
Daniel Han e7cfceadab Add linux/arm64 (DGX Spark / Grace) support via QEMU at build time
Make the docker image multi-arch so DGX Spark (GB10, sm_121, aarch64) and
the Grace-Hopper / Grace-Blackwell SoCs (GH200 arm64, GB200 arm64) pull a
natively-built arm64 child from the same manifest. Runtime emulation is
NOT involved -- QEMU is used only for the cross-compile step on x86_64
CI runners; consumers on aarch64 hosts get a normal arm64 image and CUDA
works as on any other host.

Dockerfile:
  * ARG TARGETARCH; switch unsloth extras between cu128-ampere-torch2100
    (amd64, with xformers) and huggingface (arm64, no xformers -- there
    is no cu128 aarch64 xformers wheel as of 0.0.34, so we fall back to
    Unsloth's native SDPA path; ~5-10% slowdown but functionally complete).
  * Build-time torch._C._cuda_getArchFlags() assertion: amd64 still
    requires sm_120, arm64 accepts sm_120 or sm_121.
  * Same TORCH_CUDA_ARCH_LIST on both arches; nvcc emits whatever's listed.

docker/setup_qemu.sh (new):
  One-time host setup -- registers binfmt_misc handlers via
  tonistiigi/binfmt and creates a 'unsloth-multiarch' docker-container
  buildx builder. Required only on x86_64 build hosts targeting arm64.

docker/test_locally.sh:
  --platform amd64|arm64 flag. Cross-builds verify QEMU is registered,
  then build through the in-image arch-flags assertion. Smoke + notebook
  blocks auto-skip when image arch != host arch (CUDA cannot run under
  user-space QEMU + nvidia-container-toolkit cannot bridge a QEMU guest
  to a real GPU).

.github/workflows/docker-publish.yml:
  platforms: linux/amd64,linux/arm64 (single manifest, two children).
  Timeout bumped 60 -> 150 min for the slower arm64-under-QEMU leg.
  docker/setup-qemu-action@v3 with platforms: arm64 (was implicit before).
2026-05-24 10:34:59 +00:00

59 lines
2.4 KiB
Bash
Executable file

#!/usr/bin/env bash
# One-time host setup: register QEMU binfmt handlers so `docker buildx` can
# build images for foreign architectures (e.g. linux/arm64 on an x86_64 host).
#
# After this runs once per host reboot you can do:
#
# docker buildx build --platform linux/arm64 -t unsloth-blackwell:arm64 .
# docker buildx build --platform linux/amd64,linux/arm64 --push -t YOU/img:tag .
#
# Important: QEMU is used at BUILD time only. The resulting arm64 image must
# be RUN on an aarch64 host (e.g. DGX Spark / GB10) -- CUDA does not work under
# runtime emulation. To smoke-test the arm64 image you need an actual arm64
# GPU machine.
#
# Usage:
# bash docker/setup_qemu.sh
#
# Requires: docker (28+ recommended), docker buildx plugin, root via sudo or
# membership in the `docker` group. No network access to NVIDIA registries
# is needed for this step.
set -euo pipefail
command -v docker >/dev/null || { echo "ERROR: docker not on PATH"; exit 1; }
docker buildx version >/dev/null 2>&1 || {
echo "ERROR: 'docker buildx' missing. Install:" >&2
echo " Ubuntu/Debian: sudo apt-get install -y docker-buildx" >&2
echo " RHEL/Fedora: sudo dnf install -y docker-buildx-plugin" >&2
exit 1
}
ARCH="$(uname -m)"
echo ">> host arch: ${ARCH}"
# `tonistiigi/binfmt --install all` registers handlers for every supported
# foreign arch; harmless if some are already registered. This is the canonical
# upstream Docker recipe; see https://docs.docker.com/build/building/multi-platform/
echo ">> registering QEMU binfmt handlers via tonistiigi/binfmt..."
docker run --privileged --rm tonistiigi/binfmt --install all
# Ensure we have a buildx builder that can target multiple platforms.
# The default 'docker' driver builder is single-platform; we create (or
# reuse) a 'unsloth-multiarch' container-driver builder which is multi-arch.
BUILDER="unsloth-multiarch"
if docker buildx inspect "${BUILDER}" >/dev/null 2>&1; then
echo ">> buildx builder '${BUILDER}' already exists"
else
echo ">> creating buildx builder '${BUILDER}'"
docker buildx create --name "${BUILDER}" --driver docker-container --use
fi
docker buildx use "${BUILDER}"
docker buildx inspect --bootstrap "${BUILDER}" | sed -n '1,12p'
echo
echo ">> done. Verify with:"
echo " docker buildx ls"
echo " docker buildx inspect ${BUILDER}"
echo
echo ">> cross-arch build example:"
echo " docker buildx build --platform linux/arm64 -t unsloth-blackwell:arm64 docker/"