MoE ring-of-experts: TC Pallas ring all-gather for the combine backward + x_sorted remat save - #4969

MoE ring-of-experts: TC Pallas ring all-gather for the combine backward + x_sorted remat save#4969

ultrons wants to merge 2 commits into

AI-Hypercomputer:test_964807349AI-Hypercomputer/maxtext:test_964807349from

ultrons:sivaibhav-moe-ring-collectivesultrons/maxtext:sivaibhav-moe-ring-collectivesCopy head branch name to clipboard