From scratch pretraining on english only no synthetic data, no code, 3 epochs of 1 gig of data for the ~125M param model.
Test network using
Tensor Product Attention. Other than some alterations to the attention, such as 16 heads insted of 9 and using TPA, this is the same setup as
https://huggingface.co/HuggingFaceTB/SmolLM2-135M-Instruct
One of the primary reported benefits for TPA are for inference which are not really being leveraged at all, although you can probably fit a larger bsz than traditional MHA/GQA with this. This did save about 5% on params, that amount should scale much more as the network size increases. The run time is very similar to MHA/GQA at this scale.
The code for tensor product attn is available at:
https://github.com/tensorgi/T6.