From d1e766898c62e8b85a7d94dfa9df0399c9c79721 Mon Sep 17 00:00:00 2001
From: StrangeBytesDev <dev@strangebytes.org>
Date: Tue, 24 Sep 2024 16:54:46 -0700
Subject: [PATCH] Add additional install options to Dockerfiles

Former-commit-id: 237e302b5ccf029dec0c7f859c40f122bf92e1c5
---
 docker/docker-cuda/Dockerfile         | 12 ++++++++++++
 docker/docker-cuda/docker-compose.yml |  3 +++
 docker/docker-rocm/Dockerfile         |  8 ++++++++
 docker/docker-rocm/docker-compose.yml |  2 ++
 4 files changed, 25 insertions(+)

diff --git a/docker/docker-cuda/Dockerfile b/docker/docker-cuda/Dockerfile
index a43baf21..d03ece88 100644
--- a/docker/docker-cuda/Dockerfile
+++ b/docker/docker-cuda/Dockerfile
@@ -12,6 +12,9 @@ ARG INSTALL_BNB=false
 ARG INSTALL_VLLM=false
 ARG INSTALL_DEEPSPEED=false
 ARG INSTALL_FLASHATTN=false
+ARG INSTALL_LIGER_KERNEL=false
+ARG INSTALL_HQQ=false
+ARG INSTALL_EETQ=false
 ARG PIP_INDEX=https://pypi.org/simple
 
 # Set the working directory
@@ -38,6 +41,15 @@ RUN EXTRA_PACKAGES="metrics"; \
     if [ "$INSTALL_DEEPSPEED" == "true" ]; then \
         EXTRA_PACKAGES="${EXTRA_PACKAGES},deepspeed"; \
     fi; \
+    if [ "$INSTALL_LIGER_KERNEL" == "true" ]; then \
+        EXTRA_PACKAGES="${EXTRA_PACKAGES},liger-kernel"; \
+    fi; \
+    if [ "$INSTALL_HQQ" == "true" ]; then \
+        EXTRA_PACKAGES="${EXTRA_PACKAGES},hqq"; \
+    fi; \
+    if [ "$INSTALL_EETQ" == "true" ]; then \
+        EXTRA_PACKAGES="${EXTRA_PACKAGES},eetq"; \
+    fi; \
     pip install -e ".[$EXTRA_PACKAGES]"
 
 # Rebuild flash attention
diff --git a/docker/docker-cuda/docker-compose.yml b/docker/docker-cuda/docker-compose.yml
index 16267dc3..7af06b74 100644
--- a/docker/docker-cuda/docker-compose.yml
+++ b/docker/docker-cuda/docker-compose.yml
@@ -8,6 +8,9 @@ services:
         INSTALL_VLLM: false
         INSTALL_DEEPSPEED: false
         INSTALL_FLASHATTN: false
+        INSTALL_LIGER_KERNEL: false
+        INSTALL_HQQ: false
+        INSTALL_EETQ: false
         PIP_INDEX: https://pypi.org/simple
     container_name: llamafactory
     volumes:
diff --git a/docker/docker-rocm/Dockerfile b/docker/docker-rocm/Dockerfile
index 3ce425fd..62bd78f5 100644
--- a/docker/docker-rocm/Dockerfile
+++ b/docker/docker-rocm/Dockerfile
@@ -10,6 +10,8 @@ ARG INSTALL_BNB=false
 ARG INSTALL_VLLM=false
 ARG INSTALL_DEEPSPEED=false
 ARG INSTALL_FLASHATTN=false
+ARG INSTALL_LIGER_KERNEL=false
+ARG INSTALL_HQQ=false
 ARG PIP_INDEX=https://pypi.org/simple
 
 # Set the working directory
@@ -36,6 +38,12 @@ RUN EXTRA_PACKAGES="metrics"; \
     if [ "$INSTALL_DEEPSPEED" == "true" ]; then \
         EXTRA_PACKAGES="${EXTRA_PACKAGES},deepspeed"; \
     fi; \
+    if [ "$INSTALL_LIGER_KERNEL" == "true" ]; then \
+        EXTRA_PACKAGES="${EXTRA_PACKAGES},liger-kernel"; \
+    fi; \
+    if [ "$INSTALL_HQQ" == "true" ]; then \
+        EXTRA_PACKAGES="${EXTRA_PACKAGES},hqq"; \
+    fi; \
     pip install -e ".[$EXTRA_PACKAGES]"
 
 # Rebuild flash attention
diff --git a/docker/docker-rocm/docker-compose.yml b/docker/docker-rocm/docker-compose.yml
index 923bd067..2a4ea960 100644
--- a/docker/docker-rocm/docker-compose.yml
+++ b/docker/docker-rocm/docker-compose.yml
@@ -8,6 +8,8 @@ services:
         INSTALL_VLLM: false
         INSTALL_DEEPSPEED: false
         INSTALL_FLASHATTN: false
+        INSTALL_LIGER_KERNEL: false
+        INSTALL_HQQ: false
         PIP_INDEX: https://pypi.org/simple
     container_name: llamafactory
     volumes: