DataParallel imbalanced memory usage
Simialr to this issue: https://discuss.pytorch.org/t/dataparallel-imbalanced-memory-usage/22551/12, when I run  run_lm_finetuning.py using 4 GPUs on Microsoft Azure, the first GPU will have 4000MB Memory usage while the other 3 are at 700MB. The Volatile Util for the first GPU also is at 100% while the rest are at 0%. 

It seems that the solution might be something to do with incorporating the loss calculation in the forward pass but I do not know how to solve it.
Managed to get volatile GPU to work properly but memory allocation is sitll imbalanced

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 410.78       Driver Version: 410.78       CUDA Version: 10.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla P40           Off  | 00000275:00:00.0 Off |                    0 |
| N/A   43C    P0    61W / 250W |  11151MiB / 22919MiB |     99%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla P40           Off  | 00003984:00:00.0 Off |                    0 |
| N/A   41C    P0    60W / 250W |   5979MiB / 22919MiB |     99%      Default |
+-------------------------------+----------------------+----------------------+
|   2  Tesla P40           Off  | 00005B97:00:00.0 Off |                    0 |
| N/A   38C    P0    63W / 250W |   5979MiB / 22919MiB |    100%      Default |
+-------------------------------+----------------------+----------------------+
|   3  Tesla P40           Off  | 0000EA90:00:00.0 Off |                    0 |
| N/A   42C    P0    61W / 250W |   5979MiB / 22919MiB |     99%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      5927      C   python                                     11141MiB |
|    1      5927      C   python                                      5969MiB |
|    2      5927      C   python                                      5969MiB |
|    3      5927      C   python                                      5969MiB |
+-----------------------------------------------------------------------------+<|||||>Yes, there is no mechanism to balance memory in the examples. In NVIDIA's tests, it didn't help.