build: Set triton environment variables (#370)

Signed-off-by: Will Johnson <mwjohnson728@gmail.com> Co-authored-by: Anh Uong <anh.uong@ibm.com>
foundation-model-stack · Oct 15, 2024 · 1570d04 · 1570d04
1 parent d138858
commit 1570d04
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/build/Dockerfile b/build/Dockerfile
@@ -169,6 +169,12 @@ RUN mkdir /app && \
     chown -R $USER:0 /app /tmp && \
     chmod -R g+rwX /app /tmp
 
+# Set Triton environment variables for qLoRA
+ENV TRITON_HOME="/tmp/triton_home"
+ENV TRITON_DUMP_DIR="/tmp/triton_dump_dir"
+ENV TRITON_CACHE_DIR="/tmp/triton_cache_dir"
+ENV TRITON_OVERRIDE_DIR="/tmp/triton_override_dir"
+
 # Need a better way to address these hacks
 RUN if [[ "${ENABLE_AIM}" == "true" ]] ; then \
         touch /.aim_profile && \