Minimal HuggingFace port of the bigbird-base-sparse variant of GENA-LM -- a BigBird block-sparse masked language model for long human DNA sequences (up to 4096 BPE tokens), using byte-pair (BPE) tokenization.

Vocabulary: 32,000 BPE tokens trained on DNA, including [CLS], [SEP], [PAD], [UNK], and [MASK].

The original model computes attention with DeepSpeed's block-sparse BigBird kernels. This port evaluates the same mathematical block-sparse equations in pure PyTorch: dense attention is restricted to allowed (query-block, key-block) pairs, with softmax over allowed keys only. The exact per-head checkpoint layout (master_layout) is consumed directly, using its top-left nb x nb region for nb sequence blocks. No DeepSpeed is required at inference.

The upstream wrapper forces Q/K/V to fp16 because its legacy Triton kernel accepts only fp16. This port instead evaluates the equations in the model's requested activation dtype for numerical stability. A standalone source-tensor reference is bit-exact with this mathematical implementation at all layers; exact legacy-kernel rounding/execution parity is not claimed because that kernel cannot run on Hopper (sm90).

For short sequences (<= 256 tokens / 4 blocks) the layout is fully dense, so the model behaves as an ordinary pre-LayerNorm BERT.

For short sequences (fully dense layout), all 13 representation levels and the masked-LM logits are bit-exact (max abs diff = 0.00) against the original GENA-LM code run with dense attention (no DeepSpeed), confirming weight loading, the pre-LayerNorm structure, rotary embeddings, and the block-sparse reduction. For long sequences, every layer's probabilities match an independently rebuilt masked-softmax reference exactly, are zero outside the checkpoint's master_layout, and have rows that sum to 1. Verified on GPU with PyTorch 2.7 / CUDA 12.9. This is mathematical source parity, distinct from execution parity with the unavailable legacy DeepSpeed/Triton kernel.

See the full GENA-LM collection.

flash_attention_2 is not supported for the block-sparse checkpoints because Flash Attention cannot express the checkpoint-specific arbitrary block mask. Requesting it raises an explicit error rather than silently changing the attention pattern.

Standard HuggingFace conventions. For sequence-level tasks, pool over non-padding positions or use the [CLS] token embedding as input to a prediction head.

This is a minimal, self-contained reimplementation of the GENA-LM pre-LayerNorm BigBird backbone. The block-sparse attention is computed as masked dense attention in pure PyTorch using the checkpoint's stored layout and requested activation dtype, removing the DeepSpeed dependency. The original NSP head and pooler are not included. AutoModel returns the backbone without a pooler; use the [CLS] hidden state or masked mean pooling for sequence embeddings. The input embeddings and MLM decoder are tied. Rotary cache outputs are cloned in the caller's execution mode, so inference-mode logits can safely be followed by a grad-enabled embedding forward.

Original model and code by Fishman, Kuratov, et al. (AIRI Institute). Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

MIT, following the original repository.