Scaling issue on CPU (X86 and arm) while using models from huffing face.

[choudhary-devang]

I was inferring models like (BERT, Resnet, GPT2, etc)from Hugging face on CPU and found it performing was not good compared to Pytorch, so on observing found out that models are not scaling properly in FLAX case.

benchmarking results

BERT model

cores | time in sec
1 | 38.05184221
4 | 30.66457105
8 | 29.46852469
16 | 28.98397946
32 | 28.95349479
64 | 28.77831054

System information

• OS Platform : Ubuntu 20.04
• Flax : 0.10.0, jax: 0.4.35 , jaxlib: 0.4.35
• Python version: 3.10.12
• CPU: graviton 3

Problem you have encountered:

scaling issue

script i am using:

cpu usage

[choudhary-devang]

@zaxtax, @cghawthorne, @avital, @lukaszlew can you please look into it!

[cgarciae]

@choudhary-devang you need to use jax.jit over forward_pass.