Conversation
- 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.
![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])| if isinstance(restored_opt_state, (dict, flax.core.FrozenDict)): | ||
| flat_restored = flax.traverse_util.flatten_dict(restored_opt_state) | ||
| raw_items = flat_restored.items() | ||
| else: | ||
| leaves_with_path, _ = jax.tree_util.tree_flatten_with_path(restored_opt_state) | ||
| raw_items = leaves_with_path |
There was a problem hiding this comment.
Instead of branching on isinstance(restored_opt_state, (dict, flax.core.FrozenDict)) and using flax.traverse_util.flatten_dict, we can use jax.tree_util.tree_flatten_with_path for all PyTree structures. Since _normalize_path already correctly handles both raw strings/integers and JAX KeyEntry objects (like DictKey), this simplifies the code and makes it robust to any PyTree container type.
raw_items, _ = jax.tree_util.tree_flatten_with_path(restored_opt_state)1 participant
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.