chore(customizer): bump automodel commit to r0.6.0 + add new configs

docker/Dockerfile.nhx-customizer-tasks

@@ -48,8 +48,9 @@ RUN mkdir -p /home/${USERNAME}/.cache && \
 
 # HF stack for model_spec + Nemotron configs. preserve_base_torch.txt blocks uv from
 # pulling PyPI torch (same idea as Automodel update_pyproject_pytorch.sh /
-# docker/common/uv-pytorch.toml). TRANSFORMERS_VERSION matches Automodel r0.5.0 pin.
-ARG TRANSFORMERS_VERSION=5.8.1
+# docker/common/uv-pytorch.toml). TRANSFORMERS_VERSION matches the Automodel r0.6.0 pin
+# (AUTOMODEL_COMMIT in docker/automodel/Dockerfile.nhx-automodel-base); bump them together.
+ARG TRANSFORMERS_VERSION=5.12.1
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --python ${VIRTUAL_ENV}/bin/python --no-cache \
     --overrides /app/docker/customizer/preserve_base_torch.txt \

docker/automodel/Dockerfile.nhx-automodel-base

@@ -7,17 +7,21 @@
 # PyTorch NGC base + Automodel + CUDA extension wheels for nhx-automodel images.
 # Publish target: nhx-automodel-base (slim image; builder is build-only).
 
-# Pin to the tip of the Automodel release branch (origin/r0.5.0), not main.
-# Re-pin AUTOMODEL_COMMIT when cutting to a new r0.x.y.
-ARG AUTOMODEL_COMMIT=84e85792e34a240f09194a919c1e62c9d3e47938
+# Pin to the tip of the Automodel release branch, not main.
+# Re-pin AUTOMODEL_BRANCH and AUTOMODEL_COMMIT together when cutting to a new r0.x.y, and keep
+# TRANSFORMERS_VERSION in docker/Dockerfile.nhx-customizer-tasks matching that release's
+# transformers pin.
+ARG AUTOMODEL_BRANCH=r0.6.0
+ARG AUTOMODEL_COMMIT=4066772c9c4e6f654d445ca62b9c5a3df30709ce
 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 alpine/git AS automodel-clone
+ARG AUTOMODEL_BRANCH
 ARG AUTOMODEL_COMMIT
-RUN git clone --branch main https://github.com/NVIDIA-NeMo/Automodel.git /opt/Automodel && \
+RUN git clone --branch ${AUTOMODEL_BRANCH} https://github.com/NVIDIA-NeMo/Automodel.git /opt/Automodel && \
     cd /opt/Automodel && \
     git checkout ${AUTOMODEL_COMMIT} && \
     rm -rf /opt/Automodel/.git
@@ -39,22 +43,23 @@ ENV PATH="/opt/venv/bin:/root/.local/bin:$PATH"
 RUN uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages
 
 COPY --from=automodel-clone /opt/Automodel /opt/Automodel
-COPY docker/automodel/cherry-picks /opt/cherry-picks
 
 RUN cd /opt/Automodel && \
-    bash docker/common/update_pyproject_pytorch.sh /opt/Automodel && \
-    patch -p1 < /opt/cherry-picks/3d98f6e3.diff
-
-# Sync the Automodel extras used by text training. The cherry-pick removes OpenCV
-# from base deps, including mistral-common's image extra; skipping VLM/media extras
-# keeps FFmpeg wheels out. We'll need to re-examine when Customizer supports VLM
-# and FFmpeg 8.xx only releases every 6 months
+    bash docker/common/update_pyproject_pytorch.sh /opt/Automodel
+
+# Sync the Automodel extras used by text training. Since r0.6.0, Automodel keeps
+# 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.
 RUN --mount=type=cache,target=/root/.cache/uv \
     cd /opt/Automodel && \
     UV_HTTP_TIMEOUT=120 uv sync --locked \
     --extra extra \
     --extra delta-databricks \
-    --all-groups
+    --all-groups \
+    --no-group magi
 
 # cuda extra deps not covered above (onnxscript) plus prebuilt / compiled extensions.
 RUN --mount=type=cache,target=/root/.cache/uv \
@@ -106,7 +111,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     rm -rf /opt/venv/lib/python*/site-packages/vllm \
     /opt/venv/lib/python*/site-packages/vllm-*.dist-info && \
     uv pip uninstall -y mlflow || true && \
-    rm -rf /opt/venv/lib/python*/site-packages/mlflow-3.11.1.dist-info && \
+    rm -rf /opt/venv/lib/python*/site-packages/mlflow-*.dist-info && \
     uv pip install \
     --reinstall-package mlflow-skinny \
     --reinstall-package sqlparse \

docker/automodel/README.md

@@ -84,11 +84,7 @@ Override registry: `export WHEELS_REGISTRY=...` and `export IMAGE_REGISTRY=...`
 
 **Base (`nhx-automodel-base`):** NGC PyTorch 26.05, Automodel `uv sync --locked`, pinned `transformers`/`torch`.
 
-**Automodel cherry-picks:** Platform-specific patches under `docker/automodel/cherry-picks/` are applied after `update_pyproject_pytorch.sh` and before `uv sync`. Re-pin or drop patches when upstream `r0.x.y` absorbs the same changes.
-
-| Patch | Purpose |
-|-------|---------|
-| `3d98f6e3.diff` | Drop old media deps (`decord`, `imageio-ffmpeg`) from Automodel extras, remove Automodel's unconditional `opencv-python-headless` dependency, and stop selecting `mistral-common`'s image extra for text training. The container skips Automodel's VLM extra, so FFmpeg-bearing `av` / `opencv-python-headless` wheels are not installed. |
+**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.
 
 **Customizer tasks image (`nhx-customizer-tasks`):** `uv sync --package nhx-customization-common --package nhx-models --no-dev --inexact` from the customizer workspace slice (`docker/customizer/`). Hosts shared CPU steps (`file_io`, `model_entity`, `model_spec`, LoRA sidecar) for all customization backends.