Skip to content

Latest commit

 

History

History
24 lines (21 loc) · 1.66 KB

kernels.md

File metadata and controls

24 lines (21 loc) · 1.66 KB
Raw

Setting up kernels

  • Custom CUDA layernorm kernels modified from FastFold and Oneflow accelerate about 30%-50% during different training stages. To use this feature, run the following command:

    export LAYERNORM_TYPE=fast_layernorm

    If the environment variable LAYERNORM_TYPE is set to fast_layernorm, the model will employ the layernorm we have developed; otherwise, the naive PyTorch layernorm will be adopted. The kernels will be compiled when fast_layernorm is called for the first time.

  • DeepSpeed DS4Sci_EvoformerAttention kernel is a memory-efficient attention kernel developed as part of a collaboration between OpenFold and the DeepSpeed4Science initiative. To use this feature, run the following command:

    export USE_DEEPSPEED_EVO_ATTTENTION=true

    DS4Sci_EvoformerAttention is implemented based on CUTLASS. If you use this feature, You need to clone the CUTLASS repository and specify the path to it in the environment variable CUTLASS_PATH. The Dockerfile has already include this setting:

    RUN git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git  /opt/cutlass
ENV CUTLASS_PATH=/opt/cutlass

    If you set up Protenix by pip, you can set environment variable CUTLASS_PATH as follows:

    git clone -b v3.5.1 https://github.com/NVIDIA/cutlass.git  /path/to/cutlass
export CUTLASS_PATH=/path/to/cutlass

    The kernels will be compiled when DS4Sci_EvoformerAttention is called for the first time.