More info in https://github.com/pytorch/torchtitan/pull/3199#discussion_r3177491461
More info in #3199 (comment)