diff --git a/.github/workflows/coverage.yaml b/.github/workflows/coverage.yaml
index 75415a749..df1d046e9 100644
--- a/.github/workflows/coverage.yaml
+++ b/.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:
diff --git a/build/Dockerfile b/build/Dockerfile
index 836373b54..4fec1bb52 100644
--- a/build/Dockerfile
+++ b/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