From 66238a1c9463dc689b4e4bf1121ee0d9abc366bb Mon Sep 17 00:00:00 2001
From: OlivierDehaene <23298448+OlivierDehaene@users.noreply.github.com>
Date: Mon, 11 Dec 2023 12:50:46 +0100
Subject: [PATCH] update megablocks commit

---
 Dockerfile      | 7 +++++--
 server/Makefile | 2 +-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 27c5ecee..e76c9451 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -154,6 +154,11 @@ COPY server/Makefile-vllm Makefile
 # Build specific version of vllm
 RUN make build-vllm-cuda
 
+# Build megablocks
+FROM kernel-builder as megablocks-builder
+
+RUN pip install git+https://github.com/OlivierDehaene/megablocks@16c5350f7b313a5ab52ab109feb45f159f1e5d3d
+
 # Text Generation Inference base image
 FROM nvidia/cuda:12.1.0-base-ubuntu20.04 as base
 
@@ -173,7 +178,6 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins
         ca-certificates \
         make \
         curl \
-        git \
         && rm -rf /var/lib/apt/lists/*
 
 # Copy conda with PyTorch installed
@@ -210,7 +214,6 @@ COPY server server
 COPY server/Makefile server/Makefile
 RUN cd server && \
     make gen-server && \
-    pip install git+https://github.com/OlivierDehaene/megablocks@e42a064355c540099046214bd3086ffdfe651d46 && \
     pip install -r requirements_cuda.txt && \
     pip install ".[bnb, accelerate, quantize, peft]" --no-cache-dir
 
diff --git a/server/Makefile b/server/Makefile
index 7f54aff0..86c56ee7 100644
--- a/server/Makefile
+++ b/server/Makefile
@@ -17,7 +17,7 @@ gen-server:
 	touch text_generation_server/pb/__init__.py
 
 install-megablocks:
-	pip install git+https://github.com/OlivierDehaene/megablocks@e42a064355c540099046214bd3086ffdfe651d46
+	pip install git+https://github.com/OlivierDehaene/megablocks@16c5350f7b313a5ab52ab109feb45f159f1e5d3d
 
 install: gen-server
 	pip install pip --upgrade