Fsdp & multi-host fixes

[entrpn]

• Loads params from train_state directly to save memory during inference.
• Fixes sharding issue in the conv layers of the unet/vae when running in fsdp.