feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) - #5718

feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x)#5718

khatwanimohit merged 6 commits into

atwigg/mlperfAI-Hypercomputer/maxtext:atwigg/mlperffrom

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