Flash Attention ESM2 (FAESM) This is an efficient Flash Attention implementation of ESM2 (Evolutionary Scale Modeling) that provides nearly 50% speedup and memory reduction compared to the original implementation. All the source code is from FAPLM. Give us a star if you find it useful :) Key Features Automatic Flash Attention : Automatically uses FlashAttention for up to 70% faster inference and 60% memory reduction when available Smart Fallback : Automatically falls back to PyTorch SDPA if Flash Attention is not installed Drop in Replacement : Same API as the original ESM2 models Memory Efficient : Removes padding tokens during computation for better efficiency Installation Requirements Usage One change the repo name and turn on trust remote code=True . Supported ESM Versions Model Num Layers Num Parameters fredzzp/esm2\ t36\ 3B\ UR50D 36 3B fredzzp/esm2\ t33\ 650M\ UR50D 33 650M fredzzp/esm2\ t30\ 150M\ UR50D 30 150M fredzzp/esm2\ t12\ 35M\ UR50D 12 35M fredzzp/esm2\ t6\ 8M\ UR50D 6 8M Citation If you use this implementation, please cite both the original ESM2 paper and this work: License This implementation is licensed under the MIT License. The ESM2 model weights maintain their…
We use cookies for essential functionality and analytics. You can accept or reject analytics cookies.Cookie policy