Repository: huggingface/transformers
Pull request: Update quantization CI

# What does this PR do?

This PR updates the quantization CI. We remove from testing some methods that are not used that much + we upgrade torch and cuda version. 

Changed files:
filename: docker/transformers-quantization-latest-gpu/Dockerfile
patch: @@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:12.6.0-cudnn-devel-ubuntu22.04
 LABEL maintainer="Hugging Face"
 
 ARG DEBIAN_FRONTEND=noninteractive
@@ -9,9 +9,9 @@ SHELL ["sh", "-lc"]
 # The following `ARG` are mainly used to specify the versions explicitly & directly in this docker file, and not meant
 # to be used as arguments for docker build (so far).
 
-ARG PYTORCH='2.6.0'
+ARG PYTORCH='2.8.0'
 # Example: `cu102`, `cu113`, etc.
-ARG CUDA='cu121'
+ARG CUDA='cu126'
 # Disable kernel mapping for quantization tests
 ENV DISABLE_KERNEL_MAPPING=1
 
@@ -46,50 +46,29 @@ RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/opt
 # Add PEFT
 RUN python3 -m pip install --no-cache-dir git+https://github.com/huggingface/peft@main#egg=peft
 
-# Add aqlm for quantization testing
-RUN python3 -m pip install --no-cache-dir aqlm[gpu]==1.0.2
-
-# Add vptq for quantization testing
-RUN pip install vptq
-
-# Add spqr for quantization testing
-# Commented for now as No matching distribution found we need to reach out to the authors
-# RUN python3 -m pip install --no-cache-dir spqr_quant[gpu]
-
 # Add hqq for quantization testing
 RUN python3 -m pip install --no-cache-dir hqq
 
 # For GGUF tests
 RUN python3 -m pip install --no-cache-dir gguf
 
 # Add autoawq for quantization testing
-# New release v0.2.8
 RUN python3 -m pip install --no-cache-dir autoawq[kernels]
 
 # Add quanto for quantization testing
 RUN python3 -m pip install --no-cache-dir optimum-quanto
 
-# Add eetq for quantization testing
-RUN git clone https://github.com/NetEase-FuXi/EETQ.git && cd EETQ/ && git submodule update --init --recursive && pip install .
-
-# # Add flute-kernel and fast_hadamard_transform for quantization testing
-# # Commented for now as they cause issues with the build
-# # TODO: create a new workflow to test them
-# RUN python3 -m pip install --no-cache-dir flute-kernel==0.4.1
-# RUN python3 -m pip install --no-cache-dir git+https://github.com/Dao-AILab/fast-hadamard-transform.git
-
-# Add fp-quant for quantization testing
-# Requires py3.11 but our CI runs on 3.9
-# RUN python3 -m pip install --no-cache-dir "fp-quant>=0.1.6"
-
 # Add compressed-tensors for quantization testing
 RUN python3 -m pip install --no-cache-dir compressed-tensors
 
 # Add AMD Quark for quantization testing
 RUN python3 -m pip install --no-cache-dir amd-quark
 
 # Add AutoRound for quantization testing
-RUN python3 -m pip install --no-cache-dir "auto-round>=0.5.0"
+RUN python3 -m pip install --no-cache-dir auto-round
+
+# Add torchao for quantization testing
+RUN python3 -m pip install --no-cache-dir torchao
 
 # Add transformers in editable mode
 RUN python3 -m pip install --no-cache-dir -e ./transformers[dev-torch]
@@ -103,3 +82,28 @@ RUN python3 -m pip uninstall -y flash-attn
 # When installing in editable mode, `transformers` is not recognized as a package.
 # this line must be added in order for python to be aware of transformers.
 RUN cd transformers && python3 setup.py develop
+
+# Low usage or incompatible lib, will enable later on
+
+# # Add aqlm for quantization testing
+# RUN python3 -m pip install --no-cache-dir aqlm[gpu]==1.0.2
+
+# # Add vptq for quantization testing
+# RUN pip install vptq
+
+# Add spqr for quantization testing
+# Commented for now as No matching distribution found we need to reach out to the authors
+# RUN python3 -m pip install --no-cache-dir spqr_quant[gpu]
+
+# # Add eetq for quantization testing
+# RUN git clone https://github.com/NetEase-FuXi/EETQ.git && cd EETQ/ && git submodule update --init --recursive && pip install .
+
+# # Add flute-kernel and fast_hadamard_transform for quantization testing
+# # Commented for now as they cause issues with the build
+# # TODO: create a new workflow to test them
+# RUN python3 -m pip install --no-cache-dir flute-kernel==0.4.1
+# RUN python3 -m pip install --no-cache-dir git+https://github.com/Dao-AILab/fast-hadamard-transform.git
+
+# Add fp-quant for quantization testing
+# Requires py3.11 but our CI runs on 3.9
+# RUN python3 -m pip install --no-cache-dir "fp-quant>=0.1.6"
\ No newline at end of file