improve model/optimizer loading and set directly state in sharded mode

[borisdayma]

In the training script, we do:

• load params on CPU
• load opt_state on CPU (if we restore them)
• use pjit to shard state on correct devices

Ideally we wouldn't need to go through CPU first.

Also we implemented custom abstract_init and load_on_cpu args that should be pushed to HuggingFace transformers for better support.

[borisdayma]

This has been fixed