Merge branch 'main' into fix/vision-image-token-insertion

Signed-off-by: Yashasvi Chaurasia <46622381+YashasviChaurasia@users.noreply.github.com>

Author: YashasviChaurasia

.github/workflows/image.yaml

@@ -23,7 +23,7 @@ jobs:
         run: |
             echo 'check if accelerate is installed and in the PATH'
             IMAGE_NAME=${{ vars.QUAY_REPOSITORY }}fms-hf-tuning:main-nvcr-latest
-            docker run --rm -it --entrypoint which "$IMAGE_NAME" accelerate
+            docker run --rm --entrypoint which "$IMAGE_NAME" accelerate
             echo 'checks done'
       - name: Push docker image for every commit to Quay.io as dev images
         run: |

.github/workflows/pr-command.yaml

@@ -54,9 +54,20 @@ jobs:
       - name: Build NVCR Image
         run: |
           docker build \
-            -t fms-hf-tuning:pr-${{ github.event.issue.number }}-nvcr \
+            -t ${{ vars.QUAY_REPOSITORY }}fms-hf-tuning:pr-${{ github.event.issue.number }}-nvcr \
             -f build/nvcr.Dockerfile .
 
+      - name: Login to Quay.io
+        uses: docker/login-action@v3
+        with:
+          registry: quay.io
+          username: ${{ secrets.QUAY_USERNAME }}
+          password: ${{ secrets.QUAY_ROBOT_TOKEN }}
+
+      - name: Push docker image
+        run: |
+            docker push ${{ vars.QUAY_REPOSITORY }}fms-hf-tuning:pr-${{ github.event.issue.number }}-nvcr
+
       - name: Comment build result
         if: always()
         uses: actions/github-script@v7

build/nvcr.Dockerfile

@@ -15,7 +15,7 @@
 ## Global Args #################################################################
 ## If the nvcr container is updated, ensure to check the torch and python
 ## installation version inside the dockerfile before pushing changes.
-ARG NVCR_IMAGE_VERSION=25.02-py3
+ARG NVCR_IMAGE_VERSION=25.10-py3
 
 # This is based on what is inside the NVCR image already
 ARG PYTHON_VERSION=3.12
@@ -28,58 +28,26 @@ ARG USER_UID=0
 ARG WORKDIR=/app
 ARG SOURCE_DIR=${WORKDIR}/fms-hf-tuning
 
-ARG ENABLE_FMS_ACCELERATION=true
-ARG ENABLE_AIM=false
-ARG ENABLE_MLFLOW=false
-ARG ENABLE_SCANNER=false
-ARG ENABLE_CLEARML=true
 ARG ENABLE_TRITON_KERNELS=true
-ARG ENABLE_RECOMMENDER=true
 
 # Ensures to always build mamba_ssm from source
 ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
 
-# upgrade torch as the base layer contains only torch 2.7
-RUN python -m pip install --upgrade pip && \
-    pip install --upgrade setuptools && \
-    pip install --upgrade --force-reinstall torch torchaudio torchvision --index-url https://download.pytorch.org/whl/cu128
+# install triton kernels
+RUN pip install --no-cache-dir "git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels"
 
 # Install main package + flash attention
 COPY . ${SOURCE_DIR}
 RUN cd ${SOURCE_DIR}
 
-RUN pip install --no-cache-dir ${SOURCE_DIR} && \
-    pip install --no-cache-dir --no-build-isolation ${SOURCE_DIR}[flash-attn] && \
-    pip install --no-cache-dir --no-build-isolation ${SOURCE_DIR}[mamba]
-
-# Optional extras
-RUN if [[ "${ENABLE_FMS_ACCELERATION}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[fms-accel] && \
-        python -m fms_acceleration.cli install fms_acceleration_peft && \
-        python -m fms_acceleration.cli install fms_acceleration_foak && \
-        python -m fms_acceleration.cli install fms_acceleration_aadp && \
-        python -m fms_acceleration.cli install fms_acceleration_moe && \
-        python -m fms_acceleration.cli install fms_acceleration_odm; \
-    fi
-
-RUN if [[ "${ENABLE_TRITON_KERNELS}" == "true" ]]; then \
-        pip install --no-cache-dir "git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernels"; \
-    fi
-RUN if [[ "${ENABLE_CLEARML}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[clearml]; \
-    fi
-RUN if [[ "${ENABLE_AIM}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[aim]; \
-    fi
-RUN if [[ "${ENABLE_MLFLOW}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[mlflow]; \
-    fi
-RUN if [[ "${ENABLE_SCANNER}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[scanner-dev]; \
-    fi
-RUN if [[ "${ENABLE_RECOMMENDER}" == "true" ]]; then \
-        pip install --no-cache-dir ${SOURCE_DIR}[tuning_config_recommender]; \
-    fi
+RUN pip install --no-cache-dir ${SOURCE_DIR}[flash-attn,mamba,fms-accel,clearml,tuning_config_recommender]
+
+# install fms-accel packages
+RUN python -m fms_acceleration.cli install fms_acceleration_peft && \
+    python -m fms_acceleration.cli install fms_acceleration_foak && \
+    python -m fms_acceleration.cli install fms_acceleration_aadp && \
+    python -m fms_acceleration.cli install fms_acceleration_moe && \
+    python -m fms_acceleration.cli install fms_acceleration_odm
 
 # cleanup build artifacts and caches
 RUN rm -rf /root/.cache /tmp/pip-* \
@@ -141,4 +109,6 @@ ENV TRITON_DUMP_DIR="/tmp/triton_dump_dir"
 ENV TRITON_CACHE_DIR="/tmp/triton_cache_dir"
 ENV TRITON_OVERRIDE_DIR="/tmp/triton_override_dir"
 
+RUN pip install -U accelerate
+
 CMD ["python", "/app/accelerate_launch.py"]