Fix optimizer state restoration and multi-host sharding in WAN trainer

- Fix no-op state.replace(...) when restoring opt_state and step by reassigning state.
- Add restore_optimizer_state_by_path to normalize and align Orbax checkpoint paths with live Optax PyTree structure.
- Update max_utils.device_put_replicated to convert to NumPy host arrays before callback slicing to avoid multi-host dispatch failures, and fast-path matching shardings.