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).
59 lines
2.4 KiB
Bash
Executable file
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/"
|