BERT with Flash Attention Installing dependencies To run the model on GPU, you need to install Flash Attention. You may either install from pypi (which may not work with fused dense), or from source. To install from source, clone the GitHub repository: The code provided here should work with commit 43950dd . Change to the cloned repo and install: This will compile the flash attention kernel, which will take some time. If you would like to use fused MLPs (e.g. to use activation checkpointing), you may install fused dense also from source: Configuration The config adds some new parameters: use flash attn : If True , always use flash attention. If None , use flash attention when GPU is available. If False , never use flash attention (works on CPU). window size : Size (left and right) of the local attention window. If ( 1, 1) , use global attention dense seq output : If true, we only need to pass the hidden states for the masked out token (around 15%) to the classifier heads. I set this to true for pretraining. fused mlp : Whether to use fused dense. Useful to reduce VRAM in combination with activation checkpointing mlp checkpoint lvl : One of {0, 1, 2} . Increasing this increases the…
We use cookies for essential functionality and analytics. You can accept or reject analytics cookies.Cookie policy