nroggendorff commited on
Commit
059540c
·
verified ·
1 Parent(s): 5611e6b

Update Dockerfile

Browse files
Files changed (1) hide show
  1. Dockerfile +3 -2
Dockerfile CHANGED
@@ -8,5 +8,6 @@ ENV DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
8
  TRITON_CACHE_DIR=$HOME/.cache/.triton/
9
 
10
  ENV TOKENIZERS_PARALLELISM=false \
11
- PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
12
- FLASH_ATTENTION=1
 
 
8
  TRITON_CACHE_DIR=$HOME/.cache/.triton/
9
 
10
  ENV TOKENIZERS_PARALLELISM=false \
11
+ PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128 \
12
+ FLASH_ATTENTION=1 \
13
+ ACCELERATE_USE_DEEPSPEED=true