add cuda memory and distributed metadata (#57252)
Summary:
Implementation for https://github.com/pytorch/kineto/issues/155
Pull Request resolved: https://github.com/pytorch/pytorch/pull/57252
Reviewed By: gdankel
Differential Revision: D28294662
Pulled By: ilia-cher
fbshipit-source-id: 3c71ffa333e341ff8113e891681a4905f54802dc
