Flash-BEATs

[david-rx]

Some changes for BEATs.

One functional change, adding hardcoded params into the config.

Then, a batch of changes modernizing the BEATs code to be efficient:

• Rework of attention. The previous implementation was old, manual and based on fairseq. Roughly by deleting ~500 lines of custom attention calculation and replacing with torch.sdpa(), we get a ~30-40% peak memory reduction and ~1.2x throughput depending on batch size.
• GPU-Batching of sequential spectrogram creation: BEATs hides a for-loop of spectrogram creation inside the model. It uses a niche kaldi-filterbank method which doesn't support batching, which was reimplemented with batching and GPU support -- this nearly doubled throughput.
• Removed the unnecessary float32 upcast in the GELU activation (gelu(x.float()).type_as(x) → gelu(x)) which was slowing down bf16 training a little bit.

All changes were extensively benchmarked and compared to the original implementations. Outputs of the models were close at the attention layer, after spectrgogram creation, and at the output of the final model with all changes in-place. I finally confirmed that the NatureLM training and val loss progressed identically for 10k steps after the changes.