diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index 494655d..c9d1ebe 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -56,3 +56,24 @@ jobs: labels: ${{ steps.meta.outputs.labels }} cache-from: type=gha,scope=mww-trainer-nvidia-docker cache-to: type=gha,mode=max,scope=mww-trainer-nvidia-docker + + - name: Docker metadata (Blackwell) + id: meta-blackwell + uses: docker/metadata-action@v5 + with: + images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }} + tags: | + type=raw,value=blackwell + type=ref,event=tag,suffix=-blackwell + + - name: Build and push Blackwell image + uses: docker/build-push-action@v6 + with: + context: . + file: dockerfile.blackwell + platforms: linux/amd64 + push: true + tags: ${{ steps.meta-blackwell.outputs.tags }} + labels: ${{ steps.meta-blackwell.outputs.labels }} + cache-from: type=gha,scope=mww-trainer-nvidia-docker-blackwell + cache-to: type=gha,mode=max,scope=mww-trainer-nvidia-docker-blackwell diff --git a/README.md b/README.md index 00d169c..14a99b9 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,19 @@ Tagged releases also publish matching immutable image tags: docker pull ghcr.io/tatertotterson/microwakeword:v5 ``` +RTX 50-series / Blackwell GPUs use a separate image with CUDA 12.8 and a +Python 3.13 TensorFlow build for `sm_120`: + +```bash +docker pull ghcr.io/tatertotterson/microwakeword:blackwell +docker pull ghcr.io/tatertotterson/microwakeword:v5-blackwell +``` + +Use the Blackwell image only for RTX 50-series cards. It includes the +community-built TensorFlow wheel from +[chivitiH/tensorflow-blackwell-python313](https://github.com/chivitiH/tensorflow-blackwell-python313), +which is unofficial and licensed CC BY-NC 4.0. + --- ## Run The Container @@ -39,6 +52,9 @@ docker run -d \ ``` Use a version tag such as `ghcr.io/tatertotterson/microwakeword:v5` when you want to pin a known release instead of tracking `latest`. +For RTX 50-series cards, use `ghcr.io/tatertotterson/microwakeword:blackwell` +or a pinned tag such as `ghcr.io/tatertotterson/microwakeword:v5-blackwell` +in the same `docker run` command. The flags: @@ -153,6 +169,8 @@ Personal samples are optional. Training can run with zero personal samples after Reviewed negative samples are converted into `/data/work/reviewed_negative_features/` and inserted into the training YAML as a hard-negative feature set when present. +On RTX 50-series / Blackwell GPUs, the Blackwell Docker image keeps sample generation and augmentation in the normal Python 3.12 trainer environment, then runs only the TensorFlow training/export stage in `/data/.venv-blackwell` with Python 3.13 and the Blackwell-native TensorFlow wheel. + --- ## Language Support @@ -253,3 +271,4 @@ Built on top of: - [microWakeWord](https://github.com/kahrendt/microWakeWord) - [piper-sample-generator](https://github.com/rhasspy/piper-sample-generator) +- [tensorflow-blackwell-python313](https://github.com/chivitiH/tensorflow-blackwell-python313) for the optional RTX 50-series / Blackwell image diff --git a/cli/setup_blackwell_venv b/cli/setup_blackwell_venv new file mode 100755 index 0000000..eb69f79 --- /dev/null +++ b/cli/setup_blackwell_venv @@ -0,0 +1,112 @@ +#!/bin/bash +set -euo pipefail + +PROGDIR="$(dirname "$(realpath "$0")")" +ROOTDIR="$(dirname "${PROGDIR}")" + +KNOWN_ARGS=( data-dir force python ) +source "${PROGDIR}/shell.functions" + +if [ ${#UNKNOWN_ARGS[@]} -gt 0 ] ; then + echo "Unknown argument(s): ${UNKNOWN_ARGS[*]}" >&2 + HELP=true +fi + +if [ "${HELP}" == "true" ] ; then + cat <&2 +Usage: setup_blackwell_venv [ --data-dir=/data ] [ --force ] [ --python=python3.13 ] + +Creates /data/.venv-blackwell for RTX 50 / Blackwell TensorFlow training. +Sample generation and augmentation continue to use /data/.venv. + +Environment overrides: +MWW_BLACKWELL_TF_WHEEL_URL: TensorFlow Blackwell wheel URL. + +EOF + exit 1 +fi + +[ -n "${DATA_DIR}" ] && DATA_DIR="$(realpath "${DATA_DIR}")" +[ -d "${DATA_DIR}" ] || { + echo "Data directory '${DATA_DIR}' doesn't exist." >&2 + exit 1 +} + +PYTHON="${PYTHON:-python3.13}" +VENV="${DATA_DIR}/.venv-blackwell" +MARKER="${VENV}/.mww-blackwell-venv" +TF_WHEEL_URL="${MWW_BLACKWELL_TF_WHEEL_URL:-https://github.com/chivitiH/tensorflow-blackwell-python313/releases/download/v2.22.0-selfbuilt/tensorflow-2.22.0.dev0+selfbuilt-cp313-cp313-linux_x86_64.whl}" + +if ! command -v "${PYTHON}" >/dev/null 2>&1 ; then + echo "Python 3.13 is required for the Blackwell TensorFlow wheel. Missing: ${PYTHON}" >&2 + exit 1 +fi + +if [ "${FORCE:-false}" != "true" ] && [ -x "${VENV}/bin/python" ] && [ -f "${MARKER}" ] ; then + echo " Blackwell TensorFlow venv found (skipping setup_blackwell_venv)" + exit 0 +fi + +echo "===== Setting up Blackwell TensorFlow environment ${VENV} =====" +rm -rf "${VENV}" || : +"${PYTHON}" -m venv --upgrade-deps "${VENV}" +source "${VENV}/bin/activate" + +export PIP_PROGRESS_BAR=off +export PIP_NO_COLOR=1 +export PIP_QUIET=0 + +pip_install() { + if $VERBOSE ; then + pip install "$@" || return 1 + else + { pip install "$@" || return 1 ; } | stdbuf -i0 -o0 tr -d '[:print:]' | stdbuf -i0 -o0 tr '\n' '.' + fi + echo +} + +echo " ===== Installing Blackwell TensorFlow wheel =====" +pip_install --upgrade pip setuptools wheel +pip_install "${TF_WHEEL_URL}" + +echo " ===== Installing microWakeWord training dependencies =====" +pip_install \ + audiomentations \ + audio_metadata \ + datasets \ + mmap_ninja \ + pymicro-features \ + pyyaml \ + webrtcvad-wheels \ + ai-edge-litert \ + numpy-minmax \ + numpy-rms \ + absl-py \ + "numpy==2.3.5" + +echo " ===== Checking microwakeword =====" +MWW="${DATA_DIR}/tools/microWakeWord" +if [ ! -d "${MWW}" ] || [ -n "$(git -C "${MWW}" status --porcelain 2>/dev/null || true)" ] ; then + rm -rf "${MWW}" || : + mkdir -p "${DATA_DIR}/tools" + echo " Cloning micro-wake-word to ${DATA_DIR}/tools" + git clone https://github.com/TaterTotterson/micro-wake-word "${MWW}" &>/dev/null +fi +echo " Installing microwakeword into Blackwell venv" +pip_install --no-deps -e "${MWW}" + +echo " ===== Testing Blackwell TensorFlow environment =====" +"${VENV}/bin/python" - <<'PY' +import tensorflow as tf +from ai_edge_litert.interpreter import Interpreter +from microwakeword.data import FeatureHandler +from microwakeword.inference import Model + +print("TensorFlow:", tf.__version__) +print("CUDA build:", tf.test.is_built_with_cuda()) +print("GPU:", tf.config.list_physical_devices("GPU")) +print("microWakeWord Blackwell imports available") +PY + +touch "${MARKER}" +echo "Blackwell TensorFlow environment ready: ${VENV}" diff --git a/cli/wake_word_sample_trainer b/cli/wake_word_sample_trainer index 7abf52d..9e7d17a 100644 --- a/cli/wake_word_sample_trainer +++ b/cli/wake_word_sample_trainer @@ -84,6 +84,21 @@ if [ "${IS_BLACKWELL}" = "true" ]; then echo "ℹ️ Using GPU compatibility retries; CPU fallback is ${ALLOW_CPU_FALLBACK} (override with MWW_ALLOW_CPU_FALLBACK=true|false)." fi +BLACKWELL_TF_MODE="${MWW_BLACKWELL_TF:-auto}" +BLACKWELL_TF_REQUIRED="false" +BLACKWELL_TF_ACTIVE="false" +case "${BLACKWELL_TF_MODE,,}" in + 1|true|yes|on|required) + BLACKWELL_TF_REQUIRED="true" + ;; + 0|false|no|off|disabled) + BLACKWELL_TF_MODE="disabled" + ;; + *) + BLACKWELL_TF_MODE="auto" + ;; +esac + # Enable driver-side PTX JIT fallback when ptxas/nvlink are unavailable. if [ -z "${XLA_FLAGS:-}" ]; then export XLA_FLAGS="--xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found" @@ -238,6 +253,32 @@ fi echo " Wrote training_parameters.yaml" rm -rf "${WORK_DIR}/trained_models/wakeword" +if [ "${IS_BLACKWELL}" = "true" ] && [ "${BLACKWELL_TF_MODE}" != "disabled" ]; then + BLACKWELL_SETUP="${PROGDIR}/setup_blackwell_venv" + BLACKWELL_PYTHON="${DATA_DIR}/.venv-blackwell/bin/python" + + if [ -x "${BLACKWELL_SETUP}" ] && command -v python3.13 >/dev/null 2>&1; then + echo "↪️ Preparing Blackwell-native TensorFlow training environment." + if "${BLACKWELL_SETUP}" --data-dir="${DATA_DIR}"; then + PYTHON_BIN="${BLACKWELL_PYTHON}" + BLACKWELL_TF_ACTIVE="true" + echo "✅ Blackwell TensorFlow training enabled: ${PYTHON_BIN}" + else + if [ "${BLACKWELL_TF_REQUIRED}" = "true" ]; then + echo "❌ Blackwell TensorFlow setup failed and MWW_BLACKWELL_TF is required." >&2 + exit 1 + fi + echo "⚠️ Blackwell TensorFlow setup failed; continuing with compatibility retries." + fi + else + if [ "${BLACKWELL_TF_REQUIRED}" = "true" ]; then + echo "❌ Blackwell TensorFlow was required, but python3.13/setup_blackwell_venv is unavailable." >&2 + exit 1 + fi + echo "ℹ️ Blackwell TensorFlow image support not available; continuing with compatibility retries." + fi +fi + wake_word_filename="$( echo "${WAKE_WORD}" \ | tr '[:upper:]' '[:lower:]' \ diff --git a/dockerfile.blackwell b/dockerfile.blackwell new file mode 100644 index 0000000..5562553 --- /dev/null +++ b/dockerfile.blackwell @@ -0,0 +1,54 @@ +# RTX 50 / Blackwell image +FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV CUDA_HOME=/usr/local/cuda +ENV PATH=/usr/local/cuda/bin:${PATH} +ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} +ENV MWW_BLACKWELL_IMAGE=1 +ENV MWW_BLACKWELL_TF=auto +ENV MWW_BLACKWELL_TF_WHEEL_URL=https://github.com/chivitiH/tensorflow-blackwell-python313/releases/download/v2.22.0-selfbuilt/tensorflow-2.22.0.dev0+selfbuilt-cp313-cp313-linux_x86_64.whl + +# System deps. Python 3.12 remains the main trainer/runtime venv, while +# Python 3.13 is used only for the Blackwell TensorFlow training step. +RUN apt-get update && apt-get install -y --no-install-recommends \ + software-properties-common ca-certificates curl git wget unzip patch \ + ninja-build nano less \ + && add-apt-repository -y ppa:deadsnakes/ppa \ + && apt-get update \ + && apt-get install -y --no-install-recommends \ + python3.12 python3.12-venv python3.12-dev \ + python3.13 python3.13-venv python3.13-dev \ + python3-pip python-is-python3 \ + && ldconfig \ + && rm -rf /var/lib/apt/lists/* \ + && mkdir -p /data + +# Trainer UI port +EXPOSE 8789 + +# Script root +WORKDIR /root/mww-scripts + +# Bash environment +COPY --chown=root:root --chmod=0755 .bashrc /root/ + +# Root-level entrypoints +COPY --chown=root:root --chmod=0755 \ + train_wake_word \ + run.sh \ + trainer_server.py \ + requirements.txt \ + /root/mww-scripts/ + +# CLI folder +COPY --chown=root:root cli/ /root/mww-scripts/cli/ + +# Make all CLI scripts executable (avoids "Permission denied") +RUN chmod -R a+x /root/mww-scripts/cli + +# Static UI for trainer +COPY --chown=root:root --chmod=0644 static/index.html /root/mww-scripts/static/index.html + +# trainer server +CMD ["/bin/bash", "-lc", "/root/mww-scripts/run.sh"]