[DO NOT MERGE] perf(rl_trainer_bench): opt-in STAGE_BEFORE_UPDATE step boundary (-7.2% step time on the 397B bench; stacked on #5701) - #5719

[DO NOT MERGE] perf(rl_trainer_bench): opt-in STAGE_BEFORE_UPDATE step boundary (-7.2% step time on the 397B bench; stacked on #5701)#5719

parambole wants to merge 4 commits into

atwigg/mlperfAI-Hypercomputer/maxtext:atwigg/mlperffrom

parambole/rl-bench-stage-before-updateAI-Hypercomputer/maxtext:parambole/rl-bench-stage-before-updateCopy head branch name to clipboard