Add benchmark for torch.distributed.pipeline.sync.Pipe (#49577)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49577
Repurposing the benchmarking from
https://github.com/facebookresearch/fairscale/blob/master/benchmarks/pipe.py
and pulling in a stripped down version of the benchmark into PyTorch.
Sample output:
```
Running benchmark with args: Namespace(batch_size=8, checkpoint='never', chunks=4, host='localhost', max_batch=10, num_decoder_layers=10, num_devices=4)
Number of parameters for model: 292833040
| batch 1 | wps 3593.07 | loss 25.98 | ppl 192556591553.37
| batch 2 | wps 4405.16 | loss 19.36 | ppl 256201548.33
| batch 3 | wps 4404.98 | loss 23.56 | ppl 17111244076.37
| batch 4 | wps 4413.25 | loss 27.11 | ppl 594561327825.83
| batch 5 | wps 4408.53 | loss 25.92 | ppl 181277705101.33
| batch 6 | wps 4385.64 | loss 24.92 | ppl 66592883598.50
| batch 7 | wps 4434.11 | loss 24.75 | ppl 56113635884.68
| batch 8 | wps 4441.25 | loss 24.88 | ppl 63666024212.82
| batch 9 | wps 4425.49 | loss 25.35 | ppl 101959669008.98
| batch 10 | wps 4421.05 | loss 25.34 | ppl 101597621863.94
Peak memory usage for GPUs: cuda:0: 2.38GiB, cuda:1: 3.04GiB, cuda:2: 3.04GiB, cuda:3: 3.67GiB,
```
ghstack-source-id: 118939686
Test Plan: sentinel
Reviewed By: rohan-varma
Differential Revision: D25628721
fbshipit-source-id: 41c788eed4f852aef019aec18a84cb25ad254f3a