refactor ps benchmark (#60784)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/60784
This pr refactors the ps benchmark for modular trainers.
Test Plan: Imported from OSS
Reviewed By: zou3519
Differential Revision: D29697291
Pulled By: gcramer23
fbshipit-source-id: 64579a1f5326d3cd9f32936dcf53bc243d54b71d
