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

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

copybara-service[bot] merged 1 commit into

mainAI-Hypercomputer/maxdiffusion:mainfrom

ninatu/fix_wan_optimizerAI-Hypercomputer/maxdiffusion:ninatu/fix_wan_optimizerCopy head branch name to clipboard