

TITLE
[test_model_parallelization] multiple fixes

COMMENTS
1

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 0
rocket: 0
eyes: 0

BODY
There are multiple issues with the current common `test_model_parallelization` test.

The main issue is that it uses `nvidia-smi` to take memory snapshots. This leads to 2 potential problems:
1. these tests must not be run with pytest distributed, as they rely on all the GPUs being unused - and with `-n 2` or higher it's likely to break, since `nvidia-smi` would be indiscriminately reporting memory used by other pytest workers.

    I fixed this problem first by creating a new `@require_no_pytest_distributed` decorator at  https://gist.github.com/stas00/5d58c606dbdcb82e019d6b0674f8b42a - but once the 2nd problem was fixed it no longer was needed so I removed it. I don't think we currently have any tests that must be run without `pytest-xdist`, but if any come in the future we can merge that skip decorator too.
    
2. this implementation can easily return incorrect info if CUDA device order doesn't match nvidia-smi device order (my case and this test fails for me) - so one has to use `CUDA_VISIBLE_DEVICES` to match CUDA device order to nvidia-smi's for this test to pass.

Switching to `torch.cuda.memory_allocated` fixes both problems as it measures memory usage for the current process only and in the correct order - i.e. `to(0)` always matches `memory_allocated(0)` device-wise. (the weird multi-line implementation has to do with https://github.com/pytorch/pytorch/issues/49952)

BTW, I first thought of using `pynvml`, but it would have had the same issue.` nvidia-smi` is just another front-end to `nvml`.

Other fixes:

* removes hardcoded gpt2 config
* adds `gc.collect`. One can't rely on exact memory measurements w/o manual `gc.collect` - since it gets triggered automatically at certain times as explained in its docs, which is often too late for what's being measured. Most of the time when you `del foo` it doesn't get reclaimed by `gc` right away. So the correct sequence when exact memory measurements are desired is:
   ```
            del some_variable
            gc.collect()
            torch.cuda.empty_cache()
            # now can measure memory
   ```
* last sub-test adjusted to measure against the memory snapshot before that sub-test and not at the beginning of the whole test.

`get_current_gpu_memory_use` might go into testing or benchmarking utils and perhaps need to change its name to match that it returns MBs, but it's good enough for now.

@alexorona, please let me know if it's of interest to you for the tweaks I've been proposing - please let me know if you'd like me to tag you on these.

@patrickvonplaten, @LysandreJik 