hello, and thanks!
I haven"t measured the speed difference between the EasyDeL llama and the hugging face llama cause I don"t have access to the GPUs that are required for this task I"m currently writing a white paper for that and there are some parts that this competition is not fair at all, for example, I have access to TPU-v4/3 for research and development and they are faster than the GPUs that I can have access to and PyTorch-Xla is not even close to JAX in case of using TPUs

implementation of Flash attention is edited from the beta version released by Google itself, and Ring attention is created and improved from ring-attention-near-infinite-context-length paper

But generally using flash attention is 40 % more Memory efficient and faster by ~50/40 % than flax.linen attention
Ring-attention is tested with TPUs-v4 and I have achieved 512K context length on the 6.7B model (llama) using scan_mlp, attn_mechanism ring, and shard_kv_cache set to True with sharding_mesh_array (1, 1, 6, -1)

modeling in pure Jax is much faster than flax (about 3x up to 12x faster) but in case of being user-friendly and making people able to use EasyDeL I use Jax but I have no problem with releasing a pure JAX version for every model that doesn"t seem to be a hard thing to do but when Jax is much harder than PyTorch for newcomers and flax is kinda GNN Like I stick to flax
if you have any better idea or suggest reimplementing models in pure Jax I would personally like to do that myself, but flax EasyDeL implementation of some models still seems like to be much faster than PyTorch on A10G