diff --git a/docker-bake.hcl b/docker-bake.hcl index a8ac35c99b8..7d820b5c870 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -76,6 +76,10 @@ variable "MAMBA_SSM_WHEEL_CONTEXT" { default = "" } +variable "MAGI_ATTENTION_WHEEL_CONTEXT" { + default = "" +} + variable "FFMPEG_VLM_WHEEL_CONTEXT" { default = "" } @@ -138,7 +142,7 @@ variable "RL_BASE_CONTEXT" { # The tag for base images if needed variable "WHEELS_TAG" { - default = "a696c78488b10d15d07fab80c75b42c2983c7d8c" + default = "67cc126a137e475f22c4925be21cfc76f99847f2" } variable "BAKE_CACHE_SOURCE_BRANCH" { @@ -186,6 +190,19 @@ variable "CAUSAL_CONV1D_VERSION" { default = "v1.5.3" } +# Keep in sync with the magi_attention rev in Automodel's [tool.uv.sources] +# (AUTOMODEL_COMMIT in docker/automodel/Dockerfile.nhx-automodel-base). +variable "MAGI_ATTENTION_COMMIT" { + default = "d7ea8afd44c790b65fab68a04a6a0fdd5adbf182" +} + +# nvidia-nvshmem-cu13 that magi-attention-wheel builds against and nhx-automodel-base ships; the +# wheel's RPATH points at it, so both targets take this one value. Keep in sync with the NVSHMEM +# pin in Automodel's docker/Dockerfile. +variable "NVSHMEM_VERSION" { + default = "3.6.5" +} + function "get_causal_conv1d_wheel_image" { params = [] result = "${WHEELS_REGISTRY}/causal-conv1d-wheel:${WHEELS_TAG}" @@ -196,6 +213,11 @@ function "get_mamba_ssm_wheel_image" { result = "${WHEELS_REGISTRY}/mamba-ssm-wheel:${WHEELS_TAG}" } +function "get_magi_attention_wheel_image" { + params = [] + result = "${WHEELS_REGISTRY}/magi-attention-wheel:${WHEELS_TAG}" +} + function "get_ffmpeg_vlm_wheel_image" { params = [] result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}" @@ -238,6 +260,11 @@ function "mamba_ssm_wheel_context" { result = notequal(MAMBA_SSM_WHEEL_CONTEXT, "") ? MAMBA_SSM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:mamba-ssm-wheel" : "docker-image://${get_mamba_ssm_wheel_image()}" } +function "magi_attention_wheel_context" { + params = [] + result = notequal(MAGI_ATTENTION_WHEEL_CONTEXT, "") ? MAGI_ATTENTION_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:magi-attention-wheel" : "docker-image://${get_magi_attention_wheel_image()}" +} + function "ffmpeg_vlm_wheel_context" { params = [] result = notequal(FFMPEG_VLM_WHEEL_CONTEXT, "") ? FFMPEG_VLM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:ffmpeg-vlm-wheel" : "docker-image://${get_ffmpeg_vlm_wheel_image()}" @@ -378,6 +405,7 @@ group "nhx-automodel-gpu-wheels" { "causal-conv1d-wheel", "mamba-ssm-wheel", "rl-cuda-ext-wheel", + "magi-attention-wheel", ] } @@ -840,8 +868,8 @@ target "nhx-gym-host-smoke-test" { platforms = get_platforms() } -# Python wheel builders (causal-conv1d, mamba-ssm, av, opencv-python-headless). -# CUDA extensions only ship source on PyPI; av/opencv bundle FFmpeg. Pre-built for +# Python wheel builders (causal-conv1d, mamba-ssm, magi-attention, av, opencv-python-headless). +# CUDA extensions only ship source on PyPI (or git); av/opencv bundle FFmpeg. Pre-built for # amd64 and arm64. Wheels live at /wheels/*.whl inside each image. target "causal-conv1d-wheel" { @@ -875,6 +903,21 @@ target "mamba-ssm-wheel" { platforms = get_platforms() } +target "magi-attention-wheel" { + target = "magi-attention-wheel" + context = "." + dockerfile = "docker/base/Dockerfile.python-wheels" + cache-to = maybe_registry_cache_to("magi-attention-wheel") + cache-from = maybe_registry_cache_from("magi-attention-wheel") + tags = wheel_tags("magi-attention-wheel") + output = image_output() + args = { + MAGI_ATTENTION_COMMIT = MAGI_ATTENTION_COMMIT + NVSHMEM_VERSION = NVSHMEM_VERSION + } + platforms = get_platforms() +} + target "ffmpeg-vlm-wheel" { target = "ffmpeg-vlm-wheel" context = "." @@ -1080,11 +1123,13 @@ target "nhx-automodel-base-builder" { tags = base_tags("nhx-automodel-base") output = image_output() contexts = { - causal-conv1d-wheel-image = causal_conv1d_wheel_context() - mamba-ssm-wheel-image = mamba_ssm_wheel_context() + causal-conv1d-wheel-image = causal_conv1d_wheel_context() + mamba-ssm-wheel-image = mamba_ssm_wheel_context() + magi-attention-wheel-image = magi_attention_wheel_context() } args = { NHX_COLLECT_SOURCES = NHX_COLLECT_SOURCES + NVSHMEM_VERSION = NVSHMEM_VERSION } platforms = get_platforms() } diff --git a/docker/automodel/Dockerfile.nhx-automodel-base b/docker/automodel/Dockerfile.nhx-automodel-base index dcfdffa2dce..8279afec444 100644 --- a/docker/automodel/Dockerfile.nhx-automodel-base +++ b/docker/automodel/Dockerfile.nhx-automodel-base @@ -17,6 +17,7 @@ ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.07-py3 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src +FROM magi-attention-wheel-image AS magi-attention-wheel-src FROM alpine/git AS automodel-clone ARG AUTOMODEL_BRANCH @@ -51,8 +52,8 @@ RUN cd /opt/Automodel && \ # OpenCV and FFmpeg-bearing wheels (av, imageio-ffmpeg, decord) in the opt-in # vlm-media / diffusion-media extras; skipping them keeps those wheels out. We'll need # to re-examine when Customizer supports VLM and FFmpeg 8.xx only releases every 6 months -# The magi group (MagiAttention) is a git-only CUDA build that Customizer does not -# expose yet, so skip it rather than compiling it here. +# The magi group is a git-only CUDA build; it is installed below from the prebuilt +# magi-attention-wheel image instead. RUN --mount=type=cache,target=/root/.cache/uv \ cd /opt/Automodel && \ UV_HTTP_TIMEOUT=120 uv sync --locked \ @@ -73,9 +74,11 @@ RUN --mount=type=cache,target=/root/.cache/uv \ # Install AFTER Automodel sync - uv sync drops packages not in its lockfile. RUN --mount=from=causal-conv1d-wheel-src,target=/tmp/causal-conv1d-wheel-src,readonly \ --mount=from=mamba-ssm-wheel-src,target=/tmp/mamba-ssm-wheel-src,readonly \ + --mount=from=magi-attention-wheel-src,target=/tmp/magi-attention-wheel-src,readonly \ uv pip install --no-cache-dir --no-deps \ /tmp/causal-conv1d-wheel-src/wheels/cu13.3/causal_conv1d-*cp312*.whl \ - /tmp/mamba-ssm-wheel-src/wheels/cu13.3/mamba_ssm-2.3.0-cp312*.whl + /tmp/mamba-ssm-wheel-src/wheels/cu13.3/mamba_ssm-2.3.0-cp312*.whl \ + /tmp/magi-attention-wheel-src/wheels/cu13.3/magi_attention-*cp312*.whl RUN --mount=type=cache,target=/root/.cache/uv \ git clone https://github.com/fanshiqing/grouped_gemm.git /tmp/grouped_gemm && \ @@ -178,6 +181,12 @@ RUN --mount=type=bind,source=docker/scripts/remove-listed-files.py,target=/tmp/r --mount=type=bind,source=docker/removals/files/base/pytorch-ngc-common.txt,target=/tmp/common-file-removals.txt,readonly \ python /tmp/remove-listed-files.py /tmp/common-file-removals.txt +# magi_attn_comm (magi-attention-wheel) loads libnvshmem_host through an RPATH into the system +# site-packages it was built against, so install the same NVSHMEM with the system pip, before +# PATH puts /opt/venv first. Keep NVSHMEM_VERSION equal to the magi-attention-wheel build. +ARG NVSHMEM_VERSION=3.6.5 +RUN pip install --no-cache-dir "nvidia-nvshmem-cu13==${NVSHMEM_VERSION}" + ARG NHX_COLLECT_SOURCES=0 ENV NHX_COLLECT_SOURCES=${NHX_COLLECT_SOURCES} RUN bash /usr/local/bin/collect-apt-sources.sh \ diff --git a/docker/automodel/README.md b/docker/automodel/README.md index b87f450eb6f..e47af183b4e 100644 --- a/docker/automodel/README.md +++ b/docker/automodel/README.md @@ -9,7 +9,7 @@ Two images for the **nhx-automodel** customization backend, plus the shared **`n | Image | Dockerfile | Role | |-------|------------|------| -| `nhx-automodel-base` | `docker/automodel/Dockerfile.nhx-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels | +| `nhx-automodel-base` | `docker/automodel/Dockerfile.nhx-automodel-base` | PyTorch 26.07 + Automodel + `mamba-ssm` / `causal-conv1d` / `magi-attention` wheels | | `nhx-customizer-tasks` | `docker/Dockerfile.nhx-customizer-tasks` | Shared CPU tasks (`file_io`, `model_entity`, `model_spec`, LoRA sidecar) | | `nhx-automodel-training` | `docker/automodel/Dockerfile.nhx-automodel-training` | Training (`nhx.automodel.tasks.training`) and retrieval mining (`nhx.automodel.tasks.retrieval_mine`) | @@ -27,7 +27,7 @@ Bake file: **`docker-bake.hcl`** at the NeMo Helix repo root (`context = "."`). ## Prerequisites -1. **CUDA extension wheels** (`causal-conv1d-wheel`, `mamba-ssm-wheel`) - built from this directory or pulled from NGC. The wheel Dockerfile and uv locks live under `docker/locks/` (ported from `nhx`). +1. **CUDA extension wheels** (`causal-conv1d-wheel`, `mamba-ssm-wheel`, `magi-attention-wheel`) - built from this directory or pulled from NGC. The wheel Dockerfile and uv locks live under `docker/locks/` (ported from `nhx`). 2. **Base image tag** - after building the base, set `BASE_TAG_AUTOMODEL` (or push to `BASE_REGISTRY`) before building tasks/training. @@ -82,7 +82,7 @@ Override registry: `export WHEELS_REGISTRY=...` and `export IMAGE_REGISTRY=...` ## Tasks / training runtime (platform glue) -**Base (`nhx-automodel-base`):** NGC PyTorch 26.05, Automodel `uv sync --locked`, pinned `transformers`/`torch`. +**Base (`nhx-automodel-base`):** NGC PyTorch 26.07, Automodel `uv sync --locked`, pinned `transformers`/`torch`. **Media dependencies:** The base syncs only the Automodel extras text training needs (`extra`, `delta-databricks`). Since r0.6.0, Automodel keeps `opencv-python-headless` and FFmpeg-bearing wheels (`av`, `imageio-ffmpeg`, `decord`) in the opt-in `vlm-media` / `diffusion-media` extras, so they are not installed. Earlier releases needed a local cherry-pick (`3d98f6e3.diff`) to get the same result. diff --git a/docker/base/Dockerfile.python-wheels b/docker/base/Dockerfile.python-wheels index 23cd46e075d..71662f133d5 100644 --- a/docker/base/Dockerfile.python-wheels +++ b/docker/base/Dockerfile.python-wheels @@ -8,10 +8,12 @@ # Builds Python wheels for: # - causal-conv1d (CUDA extension required by mamba-ssm) # - mamba-ssm (selective state space model) +# - magi-attention (context-parallel attention, Automodel `magi` dependency group) # - av + opencv-python-headless + decord2 (FFmpeg 8.1.2 bundled wheels for nspect / VLM) # - rl-cuda-ext-wheel (CPython 3.13 / torch 2.13.0+cu130 extensions for nhx-rl-base) # -# causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc. +# causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc; magi-attention is +# git-only and compiles CUDA kernels against the NGC torch. # av / opencv wheels bundle FFmpeg; upstream PyPI wheels lag nspect 8.1.2. # The cp311/cp312 causal-conv1d and mamba-ssm targets stay for Automodel and Unsloth. # Build those via the nhx-automodel-gpu-wheels bake group. @@ -20,6 +22,8 @@ # Build args: # CAUSAL_CONV1D_VERSION - git tag to build (default: v1.5.3) # MAMBA_22_COMMIT - git SHA or tag to build +# MAGI_ATTENTION_COMMIT - MagiAttention git SHA (match Automodel's [tool.uv.sources] rev) +# NVSHMEM_VERSION - nvidia-nvshmem-cu13 for magi-attention (match nhx-automodel-base) # CUDA_VERSION - CUDA devel image version (default: 12.8.1) # TORCH_CUDA_ARCH_LIST - semicolon-separated SM targets (default: "8.0;8.6;9.0") # FFMPEG_VERSION - FFmpeg for av/opencv wheels (default: 8.1.2) @@ -287,6 +291,46 @@ COPY --from=mamba-ssm-23-wheel-builder-py312 /wheels /wheels COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.1.1 /wheels /wheels/cu13.1.1 COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.3 /wheels /wheels/cu13.3 +# ============================================================================= +# magi-attention wheel — Python 3.12 (for Python 3.12 consumers) - Using CUDA 13.3 +# ============================================================================= +FROM mamba-wheel-base-py312-cu13.3 AS magi-attention-wheel-builder-py312-cu13.3 + +ARG MAGI_ATTENTION_COMMIT=d7ea8afd44c790b65fab68a04a6a0fdd5adbf182 +ARG NVSHMEM_VERSION=3.6.5 + +# Automodel's docker/Dockerfile installs the RDMA headers and NVSHMEM before it builds +# MagiAttention. Without NVSHMEM, magi_attn_comm still compiles its internode kernels and +# fails. MagiAttention finds NVSHMEM through the nvidia.nvshmem package and records its lib +# dir as the extension's RPATH, so nhx-automodel-base installs the same NVSHMEM_VERSION with +# the system pip, at the same path. +RUN apt-get update && \ + apt-get install -y --no-install-recommends --allow-change-held-packages rdma-core libibverbs-dev && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* && \ + pip install --no-cache-dir "nvidia-nvshmem-cu13==${NVSHMEM_VERSION}" && \ + python -c "import nvidia.nvshmem" + +# Build settings mirror Automodel's docker/Dockerfile: Hopper + Blackwell kernels only, +# FFA kernels JIT-compiled at runtime instead of prebuilt. versioningit stamps the wheel +# version from git tags (Automodel lists it as an extra build dependency). uv checks out a +# git dependency's submodules (CUTLASS, the flash-attention fork), so this clone does too. +RUN mkdir -p /wheels && \ + pip install --no-cache-dir versioningit && \ + git clone https://github.com/SandAI-org/MagiAttention.git /src/magi-attention && \ + cd /src/magi-attention && \ + git checkout ${MAGI_ATTENTION_COMMIT} && \ + git submodule update --init --recursive && \ + MAGI_ATTENTION_BUILD_COMPUTE_CAPABILITY="90,100" \ + MAGI_ATTENTION_PREBUILD_FFA=0 \ + uv build --wheel --no-build-isolation --out-dir=/wheels . && \ + rm -rf /src/magi-attention + +# The final magi-attention-wheel image contains: +# - cu13.3/magi_attention-*-cp312-*.whl (for Automodel on CUDA 13.3) +FROM scratch AS magi-attention-wheel +COPY --from=magi-attention-wheel-builder-py312-cu13.3 /wheels /wheels/cu13.3 + # ============================================================================= # FFmpeg / VLM wheels (manylinux_2_28) — av + opencv-python-headless + decord2 # ============================================================================= diff --git a/tests/smoke_gpu/test_customizer_automodel.py b/tests/smoke_gpu/test_customizer_automodel.py index 1683b892f02..35a8c39cbf2 100644 --- a/tests/smoke_gpu/test_customizer_automodel.py +++ b/tests/smoke_gpu/test_customizer_automodel.py @@ -72,6 +72,15 @@ def test_bitsandbytes_importable(): import bitsandbytes # noqa: F401 +@pytest.mark.smoke_nhx_automodel_training +def test_magi_attention_extensions_importable(): + # magi_attention's __init__ only warns when its compiled extensions fail to load, + # so import them directly to surface a missing .so or torch ABI mismatch. Loading + # magi_attn_comm also resolves libnvshmem_host through the RPATH recorded at build time, + # so this fails if the base image's NVSHMEM is missing or at another path. + from magi_attention import magi_attn_comm, magi_attn_ext # noqa: F401 + + @pytest.mark.smoke_nhx_automodel_training def test_nhx_automodel_training_importable(): import nemo_automodel # noqa: F401