Merge remote-tracking branch 'origin/main' into codex/refactor-component-builders

# Conflicts:
#	nemo_automodel/components/checkpoint/checkpointing.py
#	nemo_automodel/recipes/llm/train_ft.py
#	nemo_automodel/recipes/vlm/finetune.py
#	tests/unit_tests/recipes/test_finetune_vlm_helpers.py
#	tests/unit_tests/recipes/test_train_ft.py

Author: akoumpa

.agents/skills

@@ -59,6 +59,12 @@ jobs:
         uses: astral-sh/setup-uv@v1
         with:
           version: 0.8.22
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/local/lib/android || true
+          sudo rm -rf /usr/local/.ghcup || true
+          sudo rm -rf /usr/local/lib/node_modules || true
+          pip cache purge || true
       - name: Install ruff
         env:
           UV_PROJECT_ENVIRONMENT: ./venv
@@ -69,12 +75,6 @@ jobs:
           export PATH="./bin/:$PATH"
 
           uv sync --link-mode copy --locked --group linting
-      - name: Free up disk space
-        run: |
-          sudo rm -rf /usr/local/lib/android || true
-          sudo rm -rf /usr/local/.ghcup || true
-          sudo rm -rf /usr/local/lib/node_modules || true
-          pip cache purge || true
       - name: Run ruff
         run: |
           source ./venv/bin/activate
@@ -130,6 +130,12 @@ jobs:
         uses: astral-sh/setup-uv@v1
         with:
           version: 0.8.22
+      - name: Free up disk space
+        run: |
+          sudo rm -rf /usr/local/lib/android || true
+          sudo rm -rf /usr/local/.ghcup || true
+          sudo rm -rf /usr/local/lib/node_modules || true
+          pip cache purge || true
       - name: Install ty
         env:
           UV_PROJECT_ENVIRONMENT: ./venv

.claude/skills/skills

@@ -11,10 +11,24 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-${{ github.event.label.name || 'main' }}-${{ github.event_name }}
   cancel-in-progress: true
 
+env:
+  TORCH_CU_INDEX: https://download.pytorch.org/whl/cu130
+  CUDA_CONTAINER_IMAGE: nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04
+
 jobs:
   pre-flight:
     uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@v0.64.2
 
+  setup-install-image:
+    runs-on: ubuntu-latest
+    outputs:
+      cuda_image: ${{ steps.cfg.outputs.cuda_image }}
+    steps:
+      - id: cfg
+        shell: bash
+        run: |
+          echo "cuda_image=${CUDA_CONTAINER_IMAGE}" >> "$GITHUB_OUTPUT"
+
   pip-test:
     runs-on: ${{ matrix.arch }}
     name: Pip - Python${{ matrix.python-version }} - ${{ matrix.arch == 'ubuntu-latest' && 'AMD64/Linux' || (matrix.arch == 'ubuntu-24.04-arm' && 'ARM64/Linux' || 'ARM64/Darwin') }} - No CUDA
@@ -60,12 +74,12 @@ jobs:
   ngc-cuda-test-uv:
     runs-on: linux-amd64-cpu16
     name: UV - Python${{ matrix.python-version }} - AMD64/Linux - NGC CUDA
-    needs: [pre-flight]
+    needs: [pre-flight, setup-install-image]
     if: |
       !(needs.pre-flight.outputs.docs_only == 'true'
       || needs.pre-flight.outputs.is_deployment_workflow == 'true')
     container:
-      image: nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
+      image: ${{ needs.setup-install-image.outputs.cuda_image }}
     environment: nemo-ci
     strategy:
       fail-fast: false
@@ -142,9 +156,9 @@ jobs:
     runs-on: linux-amd64-cpu16
     name: Build CUDA wheelhouse - Python${{ matrix.python-version }} - AMD64/Linux - NGC CUDA
     container:
-      image: nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
+      image: ${{ needs.setup-install-image.outputs.cuda_image }}
     environment: nemo-ci
-    needs: [pre-flight]
+    needs: [pre-flight, setup-install-image]
     if: |
       !(needs.pre-flight.outputs.docs_only == 'true'
       || needs.pre-flight.outputs.is_deployment_workflow == 'true')
@@ -183,7 +197,7 @@ jobs:
           . ./venv/bin/activate
 
           pip install --upgrade pip
-          pip install --index-url https://download.pytorch.org/whl/cu128 "torch<=2.8.0"
+          pip install --index-url ${TORCH_CU_INDEX} "torch>=2.9.0"
           pip install numpy packaging psutil pybind11 setuptools wheel wheel_stub
 
           mkdir -p "${WHEELHOUSE_DIR}"
@@ -192,7 +206,7 @@ jobs:
             mamba-ssm \
             nv-grouped-gemm \
             transformer-engine-torch \
-            "transformer-engine[pytorch]<=2.11.0"
+            "transformer-engine[pytorch]>=2.14.1"
 
       - name: Verify CUDA wheelhouse contents
         shell: bash -x -e -u -o pipefail {0}
@@ -216,9 +230,9 @@ jobs:
     runs-on: linux-amd64-cpu16
     name: Pip - Python${{ matrix.python-version }}${{ matrix.extra-groups != '' && format('[{0}]', matrix.extra-groups) || '' }} - AMD64/Linux - NGC CUDA
     container:
-      image: nvcr.io/nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
+      image: ${{ needs.setup-install-image.outputs.cuda_image }}
     environment: nemo-ci
-    needs: [pre-flight, cuda-wheelhouse]
+    needs: [pre-flight, setup-install-image, cuda-wheelhouse]
     if: |
       !(needs.pre-flight.outputs.docs_only == 'true'
       || needs.pre-flight.outputs.is_deployment_workflow == 'true')
@@ -260,7 +274,7 @@ jobs:
           PIP_ARGS=()
 
           if [[ $EXTRA == *"fa"* || $EXTRA == *"cuda"* || $EXTRA == *"all"* ]]; then
-            pip install --find-links "${PIP_FIND_LINKS}" --index-url https://download.pytorch.org/whl/cu128 "torch<=2.8.0"
+            pip install --find-links "${PIP_FIND_LINKS}" --index-url ${TORCH_CU_INDEX} "torch>=2.9.0"
             pip install --find-links "${PIP_FIND_LINKS}" numpy packaging psutil pybind11 setuptools wheel wheel_stub
             PIP_ARGS=(--no-build-isolation)
             export TORCH_CUDA_ARCH_LIST="9.0 10.0 12.0"
@@ -335,7 +349,7 @@ jobs:
           python-binary: ./venv/bin/python
 
   install-test-summary:
-    needs: [pip-test, uv-test, ngc-cuda-test-uv, cuda-wheelhouse, ngc-cuda-test-pip, pre-flight]
+    needs: [pip-test, uv-test, ngc-cuda-test-uv, cuda-wheelhouse, ngc-cuda-test-pip, pre-flight, setup-install-image]
     runs-on: ubuntu-latest
     name: Install test summary
     if: |

.github/workflows/cicd-main.yml

@@ -35,7 +35,7 @@ on:
 
 jobs:
   code-freeze:
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_code_freeze.yml@v0.86.0
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_code_freeze.yml@v1.4.2
     with:
       library-name: NeMo-Automodel
       python-package: nemo_automodel

.github/workflows/install-test.yml

@@ -80,7 +80,7 @@ jobs:
       !cancelled() && !failure()
       && !(needs.pre-flight.outputs.docs_only == 'true'
            || needs.pre-flight.outputs.is_deployment_workflow == 'true')
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@v1.4.0
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_release_library.yml@v1.4.3
     with:
       release-ref: ${{ inputs.release-ref || github.sha }}
       python-package: nemo_automodel

.github/workflows/release-freeze.yml

@@ -18,6 +18,7 @@ on:
   pull_request:
     paths:
       - 'pyproject.toml'
+      - 'docker/common/uv-pytorch.toml'
 
 jobs:
   check_uv_lock_and_update:

.github/workflows/release.yml

@@ -21,6 +21,7 @@
 </div>
 
 ## 📣 News and Discussions
+- [06/03/2026][**Gemma 4 12B**](https://huggingface.co/google/gemma-4-12B) We now support finetuning the dense `google/gemma-4-12B` model. Check out our [recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/gemma/gemma_4_12b_hellaswag.yaml).
 - [05/27/2026][**Step-3.7-Flash**](https://huggingface.co/stepfun-ai/Step-3.7-Flash) We added model coverage for Stepfun AI's 198B-A13B MoE vision-language model, targeting image/video agentic developer workflows with a 256k context language backbone and 1.8B ViT vision tower. See the [model coverage page](https://github.com/NVIDIA-NeMo/Automodel/blob/main/docs/model-coverage/vlm/stepfun-ai/step-3-7.md).
 - [05/19/2026][**Ling 2.0**](https://huggingface.co/collections/inclusionAI/ling-20) We now support finetuning the inclusionAI Ling 2.0 MoE family (`inclusionAI/Ling-mini-2.0`, `inclusionAI/Ling-flash-2.0`, and `inclusionAI/Ling-1T`), thanks to [@Hayden727](https://github.com/Hayden727). Check out our [recipes](https://github.com/NVIDIA-NeMo/Automodel/tree/main/examples/llm_finetune/ling).
 - [05/17/2026][**ERNIE 4.5**](https://huggingface.co/baidu) and [**MiMo-V2-Flash**](https://huggingface.co/XiaomiMiMo/MiMo-V2-Flash) We now support finetuning `baidu/ERNIE-4.5-0.3B-PT`, `baidu/ERNIE-4.5-21B-A3B-PT`, and `XiaomiMiMo/MiMo-V2-Flash`. Check out our ERNIE [dense recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/ernie4_5/ernie4_5_0p3b_hellaswag.yaml), ERNIE [MoE recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/ernie4_5/ernie4_5_21b_a3b_hellaswag.yaml), and MiMo [recipe](https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/llm_finetune/mimo_v2_flash/mimo_v2_flash_hellaswag.yaml).
@@ -362,7 +363,7 @@ YAML sketch:
 checkpoint:
 enabled: true
 checkpoint_dir: ./checkpoints
-save_consolidated: true
+save_consolidated: final
 model_save_format: safetensors
 ``` -->
 

.github/workflows/sync-skills.yml

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG CUDA_IMAGE=nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
+ARG CUDA_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04
 ARG PYTORCH_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
 ARG BASE_IMAGE=cuda
 
@@ -62,8 +62,8 @@ RUN if [ -f /usr/local/bin/torchrun ]; then \
 FROM update_base_container AS automodel_dep
 
 # Install TE
-ARG INSTALL_TE=False
-ARG TE_COMMIT=release_v2.11
+ARG INSTALL_TE=True
+ARG TE_COMMIT=release_v2.14
 RUN if [ "$INSTALL_TE" = "True" ]; then \
     git clone https://github.com/NVIDIA/TransformerEngine.git && \
     cd TransformerEngine && \