diff --git a/docker-bake.hcl b/docker-bake.hcl
index ffc5529156..59a72a265c 100644
--- a/docker-bake.hcl
+++ b/docker-bake.hcl
@@ -75,12 +75,12 @@ variable "BASE_TAG_PYTHON" {
 
 # Pin for nmp-automodel-base.
 variable "BASE_TAG_AUTOMODEL" {
-  default = "d9e1851f309d3cf5389c0fc0e1049bd3c87593f8"
+  default = "65632527b258367fb8855ef431f596abb8538577"
 }
 
 # The tag for base images if needed
 variable "WHEELS_TAG" {
-  default = "c6fe97bde6a687d506a33edfee36569b8c7252f3"
+  default = "65632527b258367fb8855ef431f596abb8538577"
 }
 
 variable "BAKE_CACHE_SOURCE_BRANCH" {
diff --git a/docker/Dockerfile.nmp-automodel-base b/docker/Dockerfile.nmp-automodel-base
index 636a4b49a8..c30e2c4912 100644
--- a/docker/Dockerfile.nmp-automodel-base
+++ b/docker/Dockerfile.nmp-automodel-base
@@ -4,7 +4,10 @@
 # PyTorch NGC base + Automodel + CUDA extension wheels for nmp-automodel images.
 # Publish target: nmp-automodel-base (slim image; builder is build-only).
 
-ARG AUTOMODEL_COMMIT=0e9909f56ba48ef9761fc6f49323ba9d0a0835b2
+# 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
+ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
 
 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
 FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
@@ -16,7 +19,8 @@ RUN git clone --branch main https://github.com/NVIDIA-NeMo/Automodel.git /opt/Au
     git checkout ${AUTOMODEL_COMMIT} && \
     rm -rf /opt/Automodel/.git
 
-FROM nvcr.io/nvidia/pytorch:26.02-py3 AS nmp-automodel-base-builder
+FROM ${PYTORCH_BASE} AS nmp-automodel-base-builder
+ARG PYTORCH_BASE
 
 WORKDIR /opt
 
@@ -32,10 +36,10 @@ 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 && patch -p1 < /opt/cherry-picks/e6d2930a.diff
 
 RUN cd /opt/Automodel && \
-    bash docker/common/update_pyproject_pytorch.sh /opt/Automodel
+    bash docker/common/update_pyproject_pytorch.sh /opt/Automodel && \
+    patch -p1 < /opt/cherry-picks/3d98f6e3.diff
 
 # Sync all Automodel extras except cuda — causal-conv1d, mamba-ssm, nv-grouped-gemm,
 # and bitsandbytes are installed from prebuilt wheels / source in the steps below.
@@ -60,8 +64,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 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 \
     uv pip install --no-cache-dir --no-deps \
-    /tmp/causal-conv1d-wheel-src/wheels/cu13.1.1/causal_conv1d-*cp312*.whl \
-    /tmp/mamba-ssm-wheel-src/wheels/cu13.1.1/mamba_ssm-2.3.0-cp312*.whl
+    /tmp/causal-conv1d-wheel-src/wheels/cu13.2/causal_conv1d-*cp312*.whl \
+    /tmp/mamba-ssm-wheel-src/wheels/cu13.2/mamba_ssm-2.3.0-cp312*.whl
 
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --no-build-isolation --no-deps git+https://github.com/fanshiqing/grouped_gemm@v1.1.4
@@ -69,7 +73,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
 RUN --mount=type=cache,target=/root/.cache/uv \
     git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git && \
     cd bitsandbytes && \
-    git checkout 0.49.1 && \
+    git checkout 0.49.2 && \
     cmake -DCOMPUTE_CAPABILITY="75;80;86;87;89;90;100;103;110;120;121" -DCOMPUTE_BACKEND=cuda -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc -S . && \
     make -j"$(nproc)" && \
     uv pip install scikit-build-core --no-deps && \
@@ -87,12 +91,23 @@ 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 install \
-    "black>=26.3.1" \
-    "pyasn1>=0.6.3" \
-    "onnx>=1.21.0"
+    "aiohttp>=3.13.3,<4" \
+    "black>=26.3.1,<27" \
+    "GitPython>=3.1.50,<4" \
+    "jaraco-context>=6.1.0,<7" \
+    "jupyter-server>=2.18.0,<3" \
+    "jupyterlab>=4.5.7,<5" \
+    "mistune>=3.2.1,<4" \
+    "notebook>=7.5.6,<8" \
+    "onnx>=1.21.0,<2" \
+    "pillow>=12.2.0,<13" \
+    "pyasn1>=0.6.3,<1" \
+    "tornado>=6.5.5,<7" \
+    "urllib3>=2.7.0,<3"
 
 # Published base image (same filesystem as builder).
-FROM nvcr.io/nvidia/pytorch:26.02-py3 AS nmp-automodel-base
+FROM ${PYTORCH_BASE} AS nmp-automodel-base
+ARG PYTORCH_BASE
 COPY --from=nmp-automodel-base-builder /opt/venv /opt/venv
 COPY --from=nmp-automodel-base-builder /opt/Automodel /opt/Automodel
 # Builder pins uv 0.9.14 but does not ship it in the venv layer; PyTorch base may ship 0.10.x.
@@ -100,7 +115,8 @@ COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /bin/uv
 
 ENV VIRTUAL_ENV=/opt/venv \
     UV_PROJECT_ENVIRONMENT=/opt/venv \
-    HF_HUB_ENABLE_HF_TRANSFER=1
+    HF_HUB_ENABLE_HF_TRANSFER=1 \
+    FORCE_QWENVL_VIDEO_READER=torchcodec
 ENV PATH="/bin:/opt/venv/bin:/root/.local/bin:$PATH"
 WORKDIR /opt
 
diff --git a/docker/automodel/README.md b/docker/automodel/README.md
index 32150ec1ef..8803a6412a 100644
--- a/docker/automodel/README.md
+++ b/docker/automodel/README.md
@@ -4,7 +4,7 @@ Three images for the **nmp-automodel** customization backend. Published as flat
 
 | Image | Dockerfile | Role |
 |-------|------------|------|
-| `nmp-automodel-base` | `Dockerfile.nmp-automodel-base` | PyTorch 26.02 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
+| `nmp-automodel-base` | `Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
 | `nmp-automodel-tasks` | `Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
 | `nmp-automodel-training` | `Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
 
@@ -77,7 +77,13 @@ Override registry: `export WHEELS_REGISTRY=...` and `export IMAGE_REGISTRY=...`
 
 ## Tasks / training runtime (platform glue)
 
-**Base (`nmp-automodel-base`):** NGC PyTorch 26.02, Automodel `uv sync --locked`, pinned `transformers`/`torch`.
+**Base (`nmp-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 `decord` + `imageio-ffmpeg` (old bundled ffmpeg); use `torchcodec` for VLM video (`FORCE_QWENVL_VIDEO_READER=torchcodec`) |
 
 **Tasks image:** `uv sync --package nmp-automodel --no-dev --inexact` from the minimal workspace. CPU steps only need platform SDK glue; upgrading ancillary packages here does not affect training.
 
diff --git a/docker/automodel/cherry-picks/3d98f6e3.diff b/docker/automodel/cherry-picks/3d98f6e3.diff
new file mode 100644
index 0000000000..16227d32ff
--- /dev/null
+++ b/docker/automodel/cherry-picks/3d98f6e3.diff
@@ -0,0 +1,119 @@
+diff --git a/pyproject.toml b/pyproject.toml
+index 801ca1be..3d98f6e3 100644
+--- a/pyproject.toml
++++ b/pyproject.toml
+@@ -99,7 +99,6 @@ diffusion = [
+     "diffusers>=0.37.0",
+     "ftfy",
+     "imageio",
+-    "imageio-ffmpeg",
+     "opencv-python-headless",
+     "torchvision",
+ ]
+@@ -154,7 +153,6 @@ vlm = [
+     "pillow",
+     "qwen-omni-utils",
+     "qwen-vl-utils",
+-    "decord; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
+     "soundfile",
+     "timm<=1.0.22",
+     "torchcodec; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
+diff --git a/uv.lock b/uv.lock
+index ad6029d3..42f67bd1 100644
+--- a/uv.lock
++++ b/uv.lock
+@@ -1553,19 +1553,6 @@ wheels = [
+     { url = "https://files.pythonhosted.org/packages/4e/8c/f3147f5c4b73e7550fe5f9352eaa956ae838d5c51eb58e7a25b9f3e2643b/decorator-5.2.1-py3-none-any.whl", hash = "sha256:d316bb415a2d9e2d2b3abcc4084c6502fc09240e292cd76a76afc106a1c8e04a", size = 9190, upload-time = "2025-02-24T04:41:32.565Z" },
+ ]
+ 
+-[[package]]
+-name = "decord"
+-version = "0.6.0"
+-source = { registry = "https://pypi.org/simple" }
+-dependencies = [
+-    { name = "numpy", version = "1.26.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform != 'darwin'" },
+-    { name = "numpy", version = "2.4.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and platform_machine == 'x86_64' and sys_platform != 'darwin'" },
+-]
+-wheels = [
+-    { url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" },
+-    { url = "https://files.pythonhosted.org/packages/6c/be/e15b5b866da452e62635a7b27513f31cb581fa2ea9cc9b768b535d62a955/decord-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02665d7c4f1193a330205a791bc128f7e108eb6ae5b67144437a02f700943bad", size = 24733380, upload-time = "2021-06-14T21:30:57.766Z" },
+-]
+-
+ [[package]]
+ name = "deep-ep"
+ version = "1.2.1+7febc6e"
+@@ -2587,20 +2574,6 @@ wheels = [
+     { url = "https://files.pythonhosted.org/packages/fb/fe/301e0936b79bcab4cacc7548bf2853fc28dced0a578bab1f7ef53c9aa75b/imageio-2.37.2-py3-none-any.whl", hash = "sha256:ad9adfb20335d718c03de457358ed69f141021a333c40a53e57273d8a5bd0b9b", size = 317646, upload-time = "2025-11-04T14:29:37.948Z" },
+ ]
+ 
+-[[package]]
+-name = "imageio-ffmpeg"
+-version = "0.6.0"
+-source = { registry = "https://pypi.org/simple" }
+-sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" }
+-wheels = [
+-    { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" },
+-    { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" },
+-    { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" },
+-    { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" },
+-    { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" },
+-    { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" },
+-]
+-
+ [[package]]
+ name = "imagesize"
+ version = "1.4.1"
+@@ -3900,13 +3873,11 @@ all = [
+     { name = "boto3" },
+     { name = "causal-conv1d" },
+     { name = "databricks-sql-connector" },
+-    { name = "decord", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" },
+     { name = "deltalake" },
+     { name = "diffusers" },
+     { name = "flash-linear-attention" },
+     { name = "ftfy" },
+     { name = "imageio" },
+-    { name = "imageio-ffmpeg" },
+     { name = "mamba-ssm" },
+     { name = "mistral-common", extra = ["opencv"] },
+     { name = "nemo-run" },
+@@ -3959,7 +3930,6 @@ diffusion = [
+     { name = "diffusers" },
+     { name = "ftfy" },
+     { name = "imageio" },
+-    { name = "imageio-ffmpeg" },
+     { name = "opencv-python-headless" },
+     { name = "torchvision", version = "0.25.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" },
+     { name = "torchvision", version = "0.25.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux'" },
+@@ -3969,7 +3939,6 @@ diffusion-kernels = [
+     { name = "diffusers" },
+     { name = "ftfy" },
+     { name = "imageio" },
+-    { name = "imageio-ffmpeg" },
+     { name = "kernels" },
+     { name = "opencv-python-headless" },
+     { name = "torchvision", version = "0.25.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" },
+@@ -4007,7 +3976,6 @@ s3 = [
+ vlm = [
+     { name = "albumentations" },
+     { name = "backoff" },
+-    { name = "decord", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin'" },
+     { name = "mistral-common", extra = ["opencv"] },
+     { name = "numba", version = "0.53.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13'" },
+     { name = "numba", version = "0.62.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.13'" },
+@@ -4072,7 +4040,6 @@ requires-dist = [
+     { name = "causal-conv1d", marker = "extra == 'cuda'" },
+     { name = "databricks-sql-connector", marker = "extra == 'delta-databricks'", specifier = ">=3.0.0" },
+     { name = "datasets", specifier = ">=4.0.0" },
+-    { name = "decord", marker = "platform_machine == 'x86_64' and sys_platform != 'darwin' and extra == 'vlm'" },
+     { name = "deep-ep", marker = "extra == 'moe'", git = "https://github.com/deepseek-ai/DeepEP.git?rev=7febc6e25660af0f54d95dd781ecdcd62265ecca" },
+     { name = "deltalake", marker = "extra == 'delta-databricks'", specifier = ">=1.0.0" },
+     { name = "diffusers", marker = "extra == 'diffusion'", specifier = ">=0.37.0" },
+@@ -4081,7 +4048,6 @@ requires-dist = [
+     { name = "flashoptim", specifier = ">=0.1.3" },
+     { name = "ftfy", marker = "extra == 'diffusion'" },
+     { name = "imageio", marker = "extra == 'diffusion'" },
+-    { name = "imageio-ffmpeg", marker = "extra == 'diffusion'" },
+     { name = "kernels", marker = "extra == 'diffusion-kernels'" },
+     { name = "mamba-ssm", marker = "extra == 'cuda'" },
+     { name = "megatron-fsdp", specifier = ">=0.2.3" },
diff --git a/docker/automodel/cherry-picks/e6d2930a.diff b/docker/automodel/cherry-picks/e6d2930a.diff
deleted file mode 100644
index f15926a7d3..0000000000
--- a/docker/automodel/cherry-picks/e6d2930a.diff
+++ /dev/null
@@ -1,35 +0,0 @@
-diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py
-index ffd39951..f24b954a 100644
---- a/nemo_automodel/components/checkpoint/checkpointing.py
-+++ b/nemo_automodel/components/checkpoint/checkpointing.py
-@@ -481,8 +481,10 @@ class Checkpointer:
-         # We need to set them to False and call initialize_weights to re-initialize the weights.
-
-         # Some models cannot call initialize_weights when sharded with DTensors:
--        # - Gemma3ForConditionalGeneration: requires setting a row to zeros in the embedding matrix,
--        #   which is not supported for DTensors in the pinned torch version.
-+        # - Gemma3ForConditionalGeneration / Gemma3ForCausalLM / Phi3ForCausalLM:
-+        #   _init_weights() calls init.zeros_(module.weight[module.padding_idx]) on the
-+        #   embedding layer, which triggers DTensor redistribute and fails with sharded
-+        #   (TP) embeddings.
-         # - NemotronHForCausalLM: the HF remote code's _init_weights uses dt_bias.copy_()
-         #   which fails with DTensors. Note: v3 (MoE) has n_routed_experts and uses our custom
-         #   implementation which handles this correctly.
-@@ -491,7 +493,16 @@ class Checkpointer:
-         except Exception:
-             model_class = ""
-         is_nemotron_v2 = model_class == "NemotronHForCausalLM" and not getattr(model.config, "n_routed_experts", None)
--        skip_initialize_weights = is_nemotron_v2
-+        is_nemotron_v3_hf = (
-+            model_class == "NemotronHForCausalLM"
-+            and getattr(model.config, "n_routed_experts", None)  # is Nemotron V3
-+            and hasattr(model, "backbone")  # is HF remote code
-+        )
-+        skip_initialize_weights = (
-+            model_class in ["Gemma3ForConditionalGeneration", "Gemma3ForCausalLM", "Phi3ForCausalLM"]
-+            or is_nemotron_v2
-+            or is_nemotron_v3_hf
-+        )
-         if not skip_initialize_weights:
-             for _, module in model.named_modules():
-                 if hasattr(module, "_is_hf_initialized"):
diff --git a/docker/base/Dockerfile.mamba-wheel b/docker/base/Dockerfile.mamba-wheel
index 8c9e7198d1..c648002f35 100644
--- a/docker/base/Dockerfile.mamba-wheel
+++ b/docker/base/Dockerfile.mamba-wheel
@@ -84,7 +84,7 @@ RUN uv venv --python 3.12 /opt/venv312 && \
 ENV VIRTUAL_ENV=/opt/venv312
 ENV PATH="/opt/venv312/bin:$PATH"
 
-# special builder for 13.1.1 cuda
+# special builder for 13.1.1 cuda (unsloth / legacy 26.02 consumers)
 FROM nvcr.io/nvidia/pytorch:26.02-py3 AS mamba-wheel-base-py312-cu13.1.1
 
 ENV DEBIAN_FRONTEND=noninteractive
@@ -95,6 +95,17 @@ RUN apt-get update && apt-get install -y --no-install-recommends software-proper
 
 COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /usr/local/bin/uv
 
+# special builder for 13.2 cuda (automodel on NGC 26.05)
+FROM nvcr.io/nvidia/pytorch:26.05-py3 AS mamba-wheel-base-py312-cu13.2
+
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN apt-get update && apt-get install -y --no-install-recommends software-properties-common \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY --from=ghcr.io/astral-sh/uv:0.9.14 /uv /usr/local/bin/uv
+
 
 # =============================================================================
 # causal-conv1d wheel - Python 3.11
@@ -131,6 +142,20 @@ FROM  mamba-wheel-base-py312-cu13.1.1 AS causal-conv1d-wheel-builder-py312-cu13.
 
 ARG CAUSAL_CONV1D_VERSION=v1.5.3
 
+RUN mkdir -p /wheels && \
+    git clone --depth 1 --branch ${CAUSAL_CONV1D_VERSION} \
+    https://github.com/Dao-AILab/causal-conv1d.git /src/causal-conv1d && \
+    cd /src/causal-conv1d && \
+    CAUSAL_CONV1D_FORCE_BUILD=TRUE uv build --wheel --no-build-isolation --out-dir=/wheels . && \
+    rm -rf /src/causal-conv1d
+
+# =============================================================================
+# causal-conv1d wheel — Python 3.12 (for Python 3.12 consumers) - Using CUDA 13.2
+# =============================================================================
+FROM  mamba-wheel-base-py312-cu13.2 AS causal-conv1d-wheel-builder-py312-cu13.2
+
+ARG CAUSAL_CONV1D_VERSION=v1.5.3
+
 RUN mkdir -p /wheels && \
     git clone --depth 1 --branch ${CAUSAL_CONV1D_VERSION} \
     https://github.com/Dao-AILab/causal-conv1d.git /src/causal-conv1d && \
@@ -146,6 +171,7 @@ FROM scratch AS causal-conv1d-wheel
 COPY --from=causal-conv1d-wheel-builder /wheels /wheels
 COPY --from=causal-conv1d-wheel-builder-py312 /wheels /wheels
 COPY --from=causal-conv1d-wheel-builder-py312-cu13.1.1 /wheels /wheels/cu13.1.1
+COPY --from=causal-conv1d-wheel-builder-py312-cu13.2 /wheels /wheels/cu13.2
 
 # =============================================================================
 # mamba-ssm 2.2.5 wheel — Python 3.11 (for safe-synthesizer-tasks)
@@ -215,14 +241,33 @@ RUN mkdir -p /wheels && \
     rm -rf /src/mamba
 
 
+# =============================================================================
+# mamba-ssm 2.3.0 wheel — Python 3.12 (for Python 3.12 consumers) - Using CUDA 13.2
+# =============================================================================
+FROM mamba-wheel-base-py312-cu13.2 AS mamba-ssm-23-wheel-builder-py312-cu13.2
+
+ARG MAMBA_23_COMMIT=v2.3.0
+
+RUN mkdir -p /wheels && \
+    git clone https://github.com/state-spaces/mamba.git /src/mamba && \
+    cd /src/mamba && \
+    git checkout ${MAMBA_23_COMMIT} && \
+    sed -i "/triton/d" setup.py && \
+    sed -i "/triton/d" pyproject.toml && \
+    uv build --wheel --no-build-isolation --out-dir=/wheels . && \
+    rm -rf /src/mamba
+
+
 # The final mamba-ssm-wheel image contains:
 #   - mamba_ssm-2.2.5-cp311-*.whl  (from MAMBA_22_COMMIT=6b32be06, for Python 3.11 consumers)
 #   - mamba_ssm-2.2.5-cp312-*.whl  (from MAMBA_22_COMMIT=6b32be06, for Python 3.12 consumers)
 #   - mamba_ssm-2.3.0-cp312-*.whl  (from v2.3.0, for Python 3.12 consumers)
-#   - cu13.1.1/mamba_ssm-2.3.0-cp312-*.whl  (from v2.3.0, for Automodel on CUDA 13.1.1)
+#   - cu13.1.1/mamba_ssm-2.3.0-cp312-*.whl  (from v2.3.0, for Unsloth on CUDA 13.1.1)
+#   - cu13.2/mamba_ssm-2.3.0-cp312-*.whl  (from v2.3.0, for Automodel on CUDA 13.2)
 # Consumers must pin both version AND Python tag glob to select the correct wheel.
 FROM scratch AS mamba-ssm-wheel
 COPY --from=mamba-ssm-wheel-builder /wheels /wheels
 COPY --from=mamba-ssm-25-wheel-builder-py312 /wheels /wheels
 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.2 /wheels /wheels/cu13.2
diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/checkpoints.py b/services/automodel/src/nmp/automodel/tasks/training/backends/checkpoints.py
index 770fd73aae..16bfe8b283 100644
--- a/services/automodel/src/nmp/automodel/tasks/training/backends/checkpoints.py
+++ b/services/automodel/src/nmp/automodel/tasks/training/backends/checkpoints.py
@@ -354,7 +354,8 @@ def export_onnx(
         Path to the exported ``model.onnx`` file.
     """
     # need to import here for the tests
-    from nemo_automodel.components.models.biencoder.export_onnx import export_to_onnx
+    # llama_bidirectional is Automodel's module path; export_to_onnx uses HF AutoModel and works for any encoder checkpoint.
+    from nemo_automodel.components.models.llama_bidirectional.export_onnx import export_to_onnx
 
     logger.info(f"Exporting embedding model at path {model_path} to ONNX format at path {output_path}")
 
diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/config.py b/services/automodel/src/nmp/automodel/tasks/training/backends/config.py
index 7b7852a99f..8925b6654f 100644
--- a/services/automodel/src/nmp/automodel/tasks/training/backends/config.py
+++ b/services/automodel/src/nmp/automodel/tasks/training/backends/config.py
@@ -75,9 +75,11 @@ def compile_automodel_config(
     seed = int(os.environ.get("PL_GLOBAL_SEED", customizer_config.seed))
 
     if _is_embedding_model:
-        # Biencoder recipe creates StatefulRNG from seed value internally
-        # See: nemo_automodel/recipes/biencoder/train_biencoder.py
+        # Bi-encoder recipe creates StatefulRNG from seed internally.
+        # See: nemo_automodel/recipes/retrieval/train_bi_encoder.py
         cfg["seed"] = seed
+        # Contrastive temperature (formerly model.t in Automodel <=0.3.x).
+        cfg["temperature"] = 0.02
     else:
         # LLM recipe expects the full rng config object
         cfg["rng"] = {
@@ -103,34 +105,9 @@ def compile_automodel_config(
     if _is_embedding_model:
         cfg["model"].update(
             {
-                "_target_": "nemo_automodel.components.models.biencoder.NeMoAutoModelBiencoder.from_pretrained",
-                # Use the same encoder for both queries and passages. default value taken from Automodel example
-                "share_encoder": True,
-                # Add a trainable linear layer after pooling to reduce embedding dimension. default value taken from Automodel example
-                "add_linear_pooler": False,
-                # How to combine token embeddings into a single document/query embedding. default value taken from Automodel example
+                "_target_": "nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained",
                 "pooling": "avg",
-                # Normalize embeddings to unit length (length = 1). default value taken from Automodel example
                 "l2_normalize": True,
-                # When training an embedding model, we want it to learn that similar things should have similar embeddings
-                # and different things should have different embeddings.
-                # Temperature controls how "strict" the model is when learning these relationships.
-                # Low value (0.02), tells the model to pick the correct doc and penalizes near-misses.
-                # High value (like 1.0 that's Automodel default) tells the model to be more lenient and allows for near-misses.
-                # 0.02 is taken from the Automodel example for biencoder training.
-                "t": 0.02,
-                # Total number of passages per query during training: 1 positive + (n-1) negatives.
-                # For example, train_n_passages=5 means 1 positive and 4 negative passages per query.
-                # This differs from legacy Customizer's 'num_hard_negatives' which only counted negatives
-                # (num_hard_negatives=4 is equivalent to train_n_passages=5).
-                "train_n_passages": embedding_config.train_n_passages,
-                # Number of negative passages per query during validation.
-                "eval_negative_size": get_eval_negative_size(embedding_config),
-                # Gradient checkpointing saves memory by not storing all activations during forward pass.
-                # Instead, it recomputes them during backward pass with a memory trade-off - less memory, slower training.
-                # Useful for large models or limited GPU memory.
-                # TODO: consider exposing this in CustomizationJobInput
-                "do_gradient_checkpointing": embedding_config.do_gradient_checkpointing,
                 "use_liger_kernel": True,
                 "use_sdpa_patching": True,
             }
@@ -168,6 +145,8 @@ def compile_automodel_config(
         "ep_size": p.expert_parallel_size,
         "sequence_parallel": p.sequence_parallel,
     }
+    if _is_embedding_model and embedding_config.do_gradient_checkpointing:
+        cfg["distributed"]["activation_checkpointing"] = True
     if p.pipeline_parallel_size > 1:
         cfg["distributed"]["pipeline"] = {
             "pp_schedule": "interleaved1f1b",
@@ -679,7 +658,7 @@ def _configure_embedding_dataset(
 
     This uses retrieval_dataset_inline.make_retrieval_dataset which handles:
         - Loading inline text directly from JSONL
-        - RetrievalBiencoderCollator for tokenization and batching
+        - BiEncoderCollator for tokenization and batching
 
     Args:
         cfg: Configuration dictionary to populate.
@@ -696,14 +675,15 @@ def _configure_embedding_dataset(
         "_target_": "torchdata.stateful_dataloader.StatefulDataLoader",
         "dataset": {
             "_target_": "nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset",
+            "model_type": "bi_encoder",
             "data_dir_list": [str(train_file)],
             "data_type": "train",
-            "train_n_passages": embedding_config.train_n_passages,
+            "n_passages": embedding_config.train_n_passages,
             "seed": seed,
             "do_shuffle": True,
         },
         "collate_fn": {
-            "_target_": "nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator",
+            "_target_": "nemo_automodel.components.datasets.llm.BiEncoderCollator",
             "q_max_len": embedding_config.query_max_length,
             "p_max_len": embedding_config.passage_max_length,
             "query_prefix": embedding_config.query_prefix,
@@ -719,15 +699,16 @@ def _configure_embedding_dataset(
             "_target_": "torchdata.stateful_dataloader.StatefulDataLoader",
             "dataset": {
                 "_target_": "nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset",
+                "model_type": "bi_encoder",
                 "data_dir_list": [str(val_file)],
                 "data_type": "eval",
-                "train_n_passages": embedding_config.train_n_passages,
+                "n_passages": embedding_config.train_n_passages,
                 "eval_negative_size": get_eval_negative_size(embedding_config),
                 "seed": seed,
                 "do_shuffle": False,
             },
             "collate_fn": {
-                "_target_": "nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator",
+                "_target_": "nemo_automodel.components.datasets.llm.BiEncoderCollator",
                 "q_max_len": embedding_config.query_max_length,
                 "p_max_len": embedding_config.passage_max_length,
                 "query_prefix": embedding_config.query_prefix,
diff --git a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py
index d292255dec..609089a0d0 100644
--- a/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py
+++ b/services/automodel/src/nmp/automodel/tasks/training/backends/finetune.py
@@ -15,9 +15,9 @@
 from nemo_automodel.components.checkpoint.checkpointing import Checkpointer
 from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
 from nemo_automodel.components.training.step_scheduler import StepScheduler
-from nemo_automodel.recipes.biencoder.train_biencoder import TrainBiencoderRecipe
 from nemo_automodel.recipes.llm.kd import KnowledgeDistillationRecipeForNextTokenPrediction
 from nemo_automodel.recipes.llm.train_ft import TrainFinetuneRecipeForNextTokenPrediction
+from nemo_automodel.recipes.retrieval.train_bi_encoder import TrainBiEncoderRecipe
 from nmp.automodel.tasks.training.progress import JobsServiceProgressReporter
 from nmp.customization_common.service.context import NMPJobContext
 from nmp.customization_common.training.callbacks import TrainingProgressCallback
@@ -239,7 +239,7 @@ def create_automodel_recipe(cfg: Any) -> AutomodelRecipeWrapper:
     """Create a progress-reporting wrapper for the recipe implied by *cfg*."""
     if _is_biencoder_config(cfg):
         logger.info("Detected biencoder config, using embedding model recipe")
-        base_recipe = TrainBiencoderRecipe(cfg)
+        base_recipe = TrainBiEncoderRecipe(cfg)
     elif _is_kd_config(cfg):
         logger.info("Detected Knowledge Distillation config, using KD recipe")
         base_recipe = KnowledgeDistillationRecipeForNextTokenPrediction(cfg)
diff --git a/services/automodel/tests/contract/README.md b/services/automodel/tests/contract/README.md
index a7dc8c9fac..3868aa4831 100644
--- a/services/automodel/tests/contract/README.md
+++ b/services/automodel/tests/contract/README.md
@@ -19,8 +19,8 @@ These configs are generated by nmp-automodel's `compile_automodel_config()` —
 | meta-llama/Llama-3.2-1B-Instruct | `llama_3_2_1b_full_sft` | Full SFT | No | prompt_completion | 1 | `llm_finetune/finetune.py` |
 | meta-llama/Llama-3.2-1B-Instruct | `llama_3_2_1b_full_sft_chat` | Full SFT | No | chat | 1 | `llm_finetune/finetune.py` |
 | meta-llama/Llama-3.1-8B-Instruct | `llama_3_1_8b_full_sft_tp` | Full SFT | No | prompt_completion | 8 (TP=2) | `llm_finetune/finetune.py` |
-| nvidia/llama-nemotron-embed-1b-v2 | `embed_1b_lora` | LoRA | No | embedding | 1 | `biencoder/finetune.py` |
-| nvidia/llama-nemotron-embed-1b-v2 | `embed_1b_full_sft` | Full SFT | No | embedding | 1 | `biencoder/finetune.py` |
+| nvidia/llama-nemotron-embed-1b-v2 | `embed_1b_lora` | LoRA | No | embedding | 1 | `recipes/retrieval/train_bi_encoder.py` |
+| nvidia/llama-nemotron-embed-1b-v2 | `embed_1b_full_sft` | Full SFT | No | embedding | 1 | `recipes/retrieval/train_bi_encoder.py` |
 
 ## Sample Datasets
 
@@ -65,7 +65,7 @@ output_configs/                # flat — one YAML per input JSON
 - **Dataset paths**: Configs reference sample datasets in `./sample-datasets`. These are small (100 train / 10 val) subsets for quick validation.
 - **Sequence packing**: Packing configs have `packed_sequence.packed_sequence_size` computed from actual dataset token length statistics.
 - **Nemotron Nano**: Uses `trust_remote_code: true` since the model has custom architecture code (`NemotronHForCausalLM`).
-- **Embedding model**: Uses Automodel's biencoder recipe (`TrainBiencoderRecipe`) with retrieval dataset format, not the LLM recipe.
+- **Embedding model**: Uses Automodel's bi-encoder recipe (`TrainBiEncoderRecipe`) with retrieval dataset format, not the LLM recipe.
 - **Generation environment**: Run `generate_configs.py` in the **nmp-automodel-training** image on a GPU pod. Compiling configs requires `nemo_automodel` and its dependencies which are only available in that image.
 
 ## Running the Configs
@@ -87,7 +87,7 @@ torchrun --nproc_per_node=<NUM_GPUS> examples/llm_finetune/finetune.py --config
 ### Embedding configs
 
 ```bash
-python examples/biencoder/finetune.py --config /path/to/output_configs/<config>.yaml
+python -m nemo_automodel.recipes.retrieval.train_bi_encoder --config /path/to/output_configs/<config>.yaml
 ```
 
 ### Quick reference — all configs
@@ -113,8 +113,8 @@ python examples/llm_finetune/finetune.py --config output_configs/llama_3_2_1b_fu
 torchrun --nproc_per_node=8 examples/llm_finetune/finetune.py --config output_configs/llama_3_2_1b_full_sft_tp.yaml
 
 # --- Embedding (Biencoder) ---
-python examples/biencoder/finetune.py --config output_configs/embed_1b_lora.yaml
-python examples/biencoder/finetune.py --config output_configs/embed_1b_full_sft.yaml
+python -m nemo_automodel.recipes.retrieval.train_bi_encoder --config output_configs/embed_1b_lora.yaml
+python -m nemo_automodel.recipes.retrieval.train_bi_encoder --config output_configs/embed_1b_full_sft.yaml
 ```
 
 ## Regenerating Configs
diff --git a/services/automodel/tests/contract/output_configs/embed_1b_full_sft.yaml b/services/automodel/tests/contract/output_configs/embed_1b_full_sft.yaml
index 7f66f2a4c1..c5a44393bc 100644
--- a/services/automodel/tests/contract/output_configs/embed_1b_full_sft.yaml
+++ b/services/automodel/tests/contract/output_configs/embed_1b_full_sft.yaml
@@ -2,19 +2,14 @@ dist_env:
   backend: nccl
   timeout_minutes: 30
 seed: 1111
+temperature: 0.02
 model:
   pretrained_model_name_or_path: nvidia/llama-nemotron-embed-1b-v2
   torch_dtype: auto
   trust_remote_code: false
-  _target_: nemo_automodel.components.models.biencoder.NeMoAutoModelBiencoder.from_pretrained
-  share_encoder: true
-  add_linear_pooler: false
+  _target_: nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained
   pooling: avg
   l2_normalize: true
-  t: 0.02
-  train_n_passages: 5
-  eval_negative_size: 4
-  do_gradient_checkpointing: false
   use_liger_kernel: true
   use_sdpa_patching: true
 tokenizer:
@@ -57,14 +52,15 @@ dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   dataset:
     _target_: nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset
+    model_type: bi_encoder
     data_dir_list:
     - ./sample-datasets/embedding/training.jsonl
     data_type: train
-    train_n_passages: 5
+    n_passages: 5
     seed: 1111
     do_shuffle: true
   collate_fn:
-    _target_: nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator
+    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
     q_max_len: 512
     p_max_len: 512
     query_prefix: 'query:'
@@ -76,15 +72,16 @@ validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   dataset:
     _target_: nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset
+    model_type: bi_encoder
     data_dir_list:
     - ./sample-datasets/embedding/validation.jsonl
     data_type: eval
-    train_n_passages: 5
+    n_passages: 5
     eval_negative_size: 4
     seed: 1111
     do_shuffle: false
   collate_fn:
-    _target_: nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator
+    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
     q_max_len: 512
     p_max_len: 512
     query_prefix: 'query:'
diff --git a/services/automodel/tests/contract/output_configs/embed_1b_lora.yaml b/services/automodel/tests/contract/output_configs/embed_1b_lora.yaml
index e6eebc1c5d..48b3e4c03f 100644
--- a/services/automodel/tests/contract/output_configs/embed_1b_lora.yaml
+++ b/services/automodel/tests/contract/output_configs/embed_1b_lora.yaml
@@ -2,19 +2,14 @@ dist_env:
   backend: nccl
   timeout_minutes: 30
 seed: 1111
+temperature: 0.02
 model:
   pretrained_model_name_or_path: nvidia/llama-nemotron-embed-1b-v2
   torch_dtype: auto
   trust_remote_code: false
-  _target_: nemo_automodel.components.models.biencoder.NeMoAutoModelBiencoder.from_pretrained
-  share_encoder: true
-  add_linear_pooler: false
+  _target_: nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained
   pooling: avg
   l2_normalize: true
-  t: 0.02
-  train_n_passages: 5
-  eval_negative_size: 4
-  do_gradient_checkpointing: false
   use_liger_kernel: true
   use_sdpa_patching: true
 tokenizer:
@@ -57,14 +52,15 @@ dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   dataset:
     _target_: nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset
+    model_type: bi_encoder
     data_dir_list:
     - ./sample-datasets/embedding/training.jsonl
     data_type: train
-    train_n_passages: 5
+    n_passages: 5
     seed: 1111
     do_shuffle: true
   collate_fn:
-    _target_: nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator
+    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
     q_max_len: 512
     p_max_len: 512
     query_prefix: 'query:'
@@ -76,15 +72,16 @@ validation_dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   dataset:
     _target_: nemo_automodel.components.datasets.llm.retrieval_dataset_inline.make_retrieval_dataset
+    model_type: bi_encoder
     data_dir_list:
     - ./sample-datasets/embedding/validation.jsonl
     data_type: eval
-    train_n_passages: 5
+    n_passages: 5
     eval_negative_size: 4
     seed: 1111
     do_shuffle: false
   collate_fn:
-    _target_: nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator
+    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
     q_max_len: 512
     p_max_len: 512
     query_prefix: 'query:'