Merge branch 'main' into main

Author: dushyantbehl

.github/workflows/build-and-publish.yaml

@@ -35,6 +35,8 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+      - name: "Free up disk space"
+        uses: ./.github/actions/free-up-disk-space
       - name: Set up Python ${{ matrix.python-version.setup }}
         uses: actions/setup-python@v4
         with:

.github/workflows/coverage.yaml

@@ -10,6 +10,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - name: "Free up disk space"
+        uses: ./.github/actions/free-up-disk-space
       - name: Set up Python 3.12
         uses: actions/setup-python@v4
         with:

.github/workflows/format.yml

@@ -25,6 +25,8 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v4
+      - name: "Free up disk space"
+        uses: ./.github/actions/free-up-disk-space
       - name: Set up Python 3.12
         uses: actions/setup-python@v4
         with:

build/Dockerfile

@@ -156,9 +156,10 @@ ENV PIP_NO_BINARY=mamba-ssm,mamba_ssm
 RUN --mount=type=cache,target=/home/${USER}/.cache/pip,uid=${USER_UID} \
     python -m pip install --user wheel && \
     python -m pip install --user "$(head bdist_name)" && \
-    python -m pip install --user "$(head bdist_name)[flash-attn]" && \
     python -m pip install --user --no-build-isolation "$(head bdist_name)[mamba]"
 
+RUN python -m pip install --user --no-build-isolation "$(head bdist_name)[flash-attn]"
+
 # fms_acceleration_peft = PEFT-training, e.g., 4bit QLoRA
 # fms_acceleration_foak = Fused LoRA and triton kernels
 # fms_acceleration_aadp = Padding-Free Flash Attention Computation