# 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: ```console git clone git@github.com:Dao-AILab/flash-attention.git ``` The code provided here should work with commit `43950dd`. Change to the cloned repo and install: ```console cd flash-attention && python setup.py 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: ```console cd csrc/fused_dense_lib && python setup.py install ``` ### 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 amount of activation checkpointing within the MLP. Keep this at 0 for pretraining and use gradient accumulation instead. For embedding training, increase this as much as needed. - `last_layer_subset`: If true, we only need the compute the last layer for a subset of tokens. I left this to false. - `use_qk_norm`: Whether or not to use QK-normalization - `num_loras`: Number of LoRAs to use when initializing a `BertLoRA` model. Has no effect on other models.