Move the Mi355 to regular docker (#41989)

remi-or · web-flow · commit 80305364e2fa · 2025-11-03T16:41:06.000+01:00
* Move the Mi355 to regular docker

* Disable gfx950 compilation for FA on AMD
diff --git a/.github/workflows/self-scheduled-amd-mi355-caller.yml b/.github/workflows/self-scheduled-amd-mi355-caller.yml
@@ -21,7 +21,7 @@ jobs:
       job: run_models_gpu
       slack_report_channel: "#amd-hf-ci"
       runner_group: hfc-amd-mi355
-      docker: huggingface/testing-rocm7.0-preview
+      docker: huggingface/transformers-pytorch-amd-gpu
       ci_event: Scheduled CI (AMD) - mi355
       report_repo_id: hf-transformers-bot/transformers-ci-dummy
     secrets: inherit
@@ -33,7 +33,7 @@ jobs:
       job: run_pipelines_torch_gpu
       slack_report_channel: "#amd-hf-ci"
       runner_group: hfc-amd-mi355
-      docker: huggingface/testing-rocm7.0-preview
+      docker: huggingface/transformers-pytorch-amd-gpu
       ci_event: Scheduled CI (AMD) - mi355
       report_repo_id: hf-transformers-bot/transformers-ci-dummy
     secrets: inherit
@@ -45,7 +45,7 @@ jobs:
       job: run_examples_gpu
       slack_report_channel: "#amd-hf-ci"
       runner_group: hfc-amd-mi355
-      docker: huggingface/testing-rocm7.0-preview
+      docker: huggingface/transformers-pytorch-amd-gpu
       ci_event: Scheduled CI (AMD) - mi355
       report_repo_id: hf-transformers-bot/transformers-ci-dummy
     secrets: inherit
diff --git a/docker/transformers-pytorch-amd-gpu/Dockerfile b/docker/transformers-pytorch-amd-gpu/Dockerfile
@@ -39,7 +39,7 @@ RUN python3 -m pip install --no-cache-dir "torchcodec==0.5"
 # Install flash attention from source. Tested with commit 6387433156558135a998d5568a9d74c1778666d8
 RUN git clone https://github.com/ROCm/flash-attention/ -b tridao && \
     cd flash-attention && \
-    GPU_ARCHS="gfx942;gfx950" python setup.py install  
-# GPU_ARCHS builds for MI300, MI325 and MI355
+    GPU_ARCHS="gfx942" python setup.py install  
+# GPU_ARCHS builds for MI300, MI325 but not MI355: we would need to add `;gfx950` but it takes too long to build.
 
 RUN python3 -m pip install --no-cache-dir einops
