Repository: pytorch
Issue: [pipeline] gpu util + peak mem reporting to tune partitions and chunks 
Labels: oncall: distributed,feature,triaged,pipeline parallelism

## 🚀 Feature
<!-- A clear and concise description of the feature proposal -->

This is wrt the new Pipeline feature in pytorch master.

In order to optimize partitioning and the `chunks` hyper-parameter one needs to get solid stats on:
1. gpu utilization % - we want as close as possible to 100% for each gpu - this is primarily for the `chunks` arg tuneup
2. peak memory - we want it to be as balanced as possible across all gpus - this is primarily for optimizing the layer partitioning

It'd be good if pipeline came with ready to use tools:
- manually - important if multiple Pipes are used so that the user will have to take the measurements where they need to - i.e. before pipes (explanation: in the current experiment I converted 1 model to use 2 pipes)
- part of `Pipe` API, may be if an argument is passed to activate reporting in its init.

but manual approach would be great to start with, especially while we experiment and test what would really help for the best pipeline setup. Also because the application may have a lot more going on before the pipeline which would also use more gpu, so reporting gpu usage at the end of pipeline could lead to misleading impression - i.e. the actual peak memory usage could be bigger than the peak that occurred in the course of running the pipeline.

## peak memory

For (2) peak memory - I tried this:
```
    def pipeline_finalize(self):
        import torch
        import json
        n_gpus = torch.cuda.device_count()
        mem_map = {}
        for id in range(n_gpus):
            with torch.cuda.device(id):
                # XXX: this doesn't seem to report the right thing - getting much lower numbers
                mem_map[id] = torch.cuda.max_memory_allocated() >> 20

        logger.info(f"peak memory usage per device in MBs:\n{json.dumps(mem_map, sort_keys=True, indent=4)}")
        # reset for the next train/eval/predict stage
        # XXX: probably should do in the trainer?
        torch.cuda.reset_peak_memory_stats()
```

but I seem to get much lower numbers than what I see in nvidia-smi. Any ideas why? I'm aware of caching and cuda kernels overhead - it's still much lower. I was even concerned that perhaps threads weren't sharing the same per-process peak memory usage counter, but I verified they do.

Since gpu 0 is often busy doing a lot more than pipeline, often to achieve a good balance one needs to put less layers on gpu 0.

## gpu utilization %

for (1) gpu utilization % - I don't think pytorch measures gpu utilization - probably need to use `pynvml` and a thread that will monitor the % and report average and peak numbers per gpu. Of course, I'm open to any other ideas.

To clarify - if I split the model over 2 gpus and I get a total of 100% utilization over 2 gpus (out of 200%) - this is terrible - it means the pipeline isn't doing anything good over the naive implementation, where all but one gpu works at any given time. That's of course, if the same model setup w/o pipeline) can saturate a single gpu to 100% utilization. We can't expect 200% utilization due to copying, but it should be over 100% for the pipeline to make sense over naive MP.

If I understand how it works correctly - nvidia-smi only reports gpu util % when kernels are doing something, so data copying is most likely not included. Which is easy to see if the trainer is not fast enough loading data onto the gpu.

Thank you.

 @pritamdamania87


cc @osalpekar @jiayisuse @lw @beauby @pritamdamania87 @mrshenli @jjlilley @gqchen @rohan-varma @pietern @zhaojuanmao @satgera @aazzolini @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23