Question about loading weights with orbax

[yuanfz98]

Hello,

I wonder if it's normal that loading weights of unquantized dsr1 (~801GB) takes more than 30min. Is there any method to speed it up ?

I noticed that the project has not been updated since 2 months. Unlike some heavy projects that englobe too many models, personnally I think it is a valuable project providing a simple but performant way to implement LLM Models with JAX. I wish I can see new contributions.