Repository navigation
Fix optimizer state restoration and multi-host sharding in WAN trainer - #484
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors optimizer state restoration in base_wan_trainer.py by matching parameter paths and applying sharding, improves checkpoint saving logic, and optimizes array replication in max_utils.py. Feedback recommends avoiding synchronous device-to-host transfers in max_utils.py by using jax.device_put directly on JAX arrays, and simplifying the PyTree flattening logic in base_wan_trainer.py by using jax.tree_util.tree_flatten_with_path universally.
| arr = getattr(x, "value", x) | ||
| arr_np = np.asarray(arr) | ||
| shd = getattr(sharding, "value", sharding) | ||
| res = jax.make_array_from_callback(arr.shape, shd, lambda index: arr[index]) | ||
| res = jax.make_array_from_callback(arr_np.shape, shd, lambda index: arr_np[index]) |
There was a problem hiding this comment.
Calling np.asarray(arr) on a jax.Array triggers a synchronous device-to-host transfer, which blocks the CPU and degrades performance. Instead, if arr is already a jax.Array, we can directly use jax.device_put(arr, shd) to perform an efficient in-device or cross-device/host resharding without copying the data back to the host CPU.
arr = getattr(x, \"value\", x)\n shd = getattr(sharding, \"value\", sharding)\n if isinstance(arr, jax.Array):\n res = jax.device_put(arr, shd)\n else:\n arr_np = np.asarray(arr)\n res = jax.make_array_from_callback(arr_np.shape, shd, lambda index: arr_np[index])There was a problem hiding this comment.
For Wan optimizer restoration, jax.device_put(arr, shd) does not work in multi-host setups and fails with:
ValueError: When the second argument to device_put is a Device, the first argument must be a fully addressable array or a non-addressable array with a single device sharding. Got value with devices {CpuDevice(id=4096), CpuDevice(id=6144), CpuDevice(id=2048), CpuDevice(id=0)}
Additionally, the array is already on the CPU host: the restored checkpoint leaves are loaded from Orbax on a replicated CPU mesh, not on accelerator devices, so np.asarray(arr) is not performing an expensive device-to-host transfer.
09849ce to
0bc77cd
Compare
|
@Perseus14 can you take a look? |
1db140a to
bb66331
Compare
0ae4965 to
79810ea
Compare
- 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.
79810ea to
b8a3fba
Compare
3 participants
Previously, optimizer state restoration was silently failing as a no-op because
state.replace(opt_state=...)returns a new TrainState rather than mutating in place, and its return value was discarded.Fixing the assignment exposed that the restored optimizer state had no device sharding applied:
Restored optimizer state from checkpoints must be explicitly sharded across the device mesh (
state_shardings.opt_state) to run on device.Because Orbax checkpoint key paths diverge from live Optax PyTree structures (string indices vs. JAX KeyEntry objects and trailing
.valuewrappers),restore_optimizer_state_by_pathwas added to map restored checkpoint leaves onto the target sharded PyTree structure.In
max_utils.device_put_replicated, populating device shards from host arrays viajax.make_array_from_callbackcaused eager cross-process coordination failures across multi-host setups. Converting input arrays to local NumPy arrays (np.asarray(arr)) allows clean, process-local slicing per device shard without inter-host dispatch.