update torch

This commit is contained in:
Mohit Sharma 2024-08-05 12:47:21 +00:00
parent 0ad78d20a5
commit 55e6059eb1
2 changed files with 19 additions and 18 deletions

View File

@ -98,24 +98,26 @@ RUN pip uninstall -y triton && \
cd triton/python && \ cd triton/python && \
pip install . pip install .
RUN git clone --depth 1 --recursive --single-branch --branch 2.3-patched https://github.com/fxmarty/pytorch.git pytorch && cd pytorch && pip install -r requirements.txt --no-cache-dir # RUN git clone --depth 1 --recursive --single-branch --branch 2.3-patched https://github.com/fxmarty/pytorch.git pytorch && cd pytorch && pip install -r requirements.txt --no-cache-dir
ARG _GLIBCXX_USE_CXX11_ABI="1" # ARG _GLIBCXX_USE_CXX11_ABI="1"
ARG CMAKE_PREFIX_PATH="/opt/conda" # ARG CMAKE_PREFIX_PATH="/opt/conda"
ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942" # ARG PYTORCH_ROCM_ARCH="gfx90a;gfx942"
ARG BUILD_CAFFE2="0" \ # ARG BUILD_CAFFE2="0" \
BUILD_CAFFE2_OPS="0" \ # BUILD_CAFFE2_OPS="0" \
USE_CUDA="0" \ # USE_CUDA="0" \
USE_ROCM="1" \ # USE_ROCM="1" \
BUILD_TEST="0" \ # BUILD_TEST="0" \
USE_FBGEMM="0" \ # USE_FBGEMM="0" \
USE_NNPACK="0" \ # USE_NNPACK="0" \
USE_QNNPACK="0" \ # USE_QNNPACK="0" \
USE_XNNPACK="0" \ # USE_XNNPACK="0" \
USE_FLASH_ATTENTION="1" \ # USE_FLASH_ATTENTION="1" \
USE_MEM_EFF_ATTENTION="0" # USE_MEM_EFF_ATTENTION="0"
RUN cd pytorch && python tools/amd_build/build_amd.py && python setup.py install # RUN cd pytorch && python tools/amd_build/build_amd.py && python setup.py install
RUN pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.1
# Set AS recommended: https://github.com/ROCm/triton/wiki/A-script-to-set-program-execution-environment-in-ROCm # Set AS recommended: https://github.com/ROCm/triton/wiki/A-script-to-set-program-execution-environment-in-ROCm
ENV HIP_FORCE_DEV_KERNARG=1 ENV HIP_FORCE_DEV_KERNARG=1

View File

@ -1150,8 +1150,7 @@ class FlashCausalLM(Model):
elif CUDA_GRAPHS is not None: elif CUDA_GRAPHS is not None:
tuning_sequences = CUDA_GRAPHS tuning_sequences = CUDA_GRAPHS
else: else:
# For seqlen = 1, we dispatch to LLMM1 kernel. tuning_sequences = [1, 2, 3, 4, 5, 6, 7]
tuning_sequences = [2, 3, 4, 5, 6, 7]
tunableop_filepath = os.path.join( tunableop_filepath = os.path.join(
HUGGINGFACE_HUB_CACHE, HUGGINGFACE_HUB_CACHE,