support teraio, based on torchtitan github commit ac90c36e39c6274f9beaf76922627665b6553905
