Fix the gradients' device memory space when enabling parameter offloading#3503

copybara-service[bot] merged 1 commit intoAI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

zhenying-liu:param-offload-fixzhenying-liu/jane-maxtext:param-offload-fixCopy head branch name to clipboard