I developed a custom GPU kernel in Triton to accelerate training with a sparse feed-forward network (FFN).

With K=256, activating 256 of 4,096 FFN neurons per token (6.25%), the sparse model achieved 36% higher training throughput than its dense counterpart.

The kernel has so far been tested on a single NVIDIA RTX PRO 6000 GPU and requires GPU hardware with BF16 support.

Whether a model trained with a sparse FFN and densified afterwards can match the quality of a model trained densely remains an open question. These early results have motivated me to continue training a larger model using this approach.

More about EUGEN and my work: https://jakubcwirlej.com/ai-dev/