[Train] Cuda OOM on shutdown

[amogkam]

What happened + What you expected to happen

On TorchTrainer shutdown, we call torch.cuda.empty_cache(). However, this operation writes data to GPU memory, and by default uses device 0: https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898/2.

When training on instances with many GPUs, this can lead to OOM.

:actor_name:TorchTrainer
[ERROR 2022-11-30 21:00:25,846] worker.py: 400  Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): �[36mray::RayTrainWorker._RayTrainWorker__execute()�[39m (pid=9616, ip=172.31.5.55, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x7f3708f1d9d0>)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 31, in __execute
    raise skipped from exception_cause(skipped)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 28, in __execute
    return func(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/torch/config.py", line 118, in _shutdown_torch
    torch.cuda.empty_cache()
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/torch/cuda/memory.py", line 114, in empty_cache
    torch._C._cuda_emptyCache()
RuntimeError: CUDA error: out of memory
[ERROR 2022-11-30 21:00:25,847] worker.py: 400  Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): �[36mray::RayTrainWorker._RayTrainWorker__execute()�[39m (pid=9613, ip=172.31.5.55, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x7fa6390c39a0>)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 31, in __execute
    raise skipped from exception_cause(skipped)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 28, in __execute
    return func(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/torch/config.py", line 118, in _shutdown_torch
    torch.cuda.empty_cache()
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/torch/cuda/memory.py", line 114, in empty_cache
    torch._C._cuda_emptyCache()
RuntimeError: CUDA error: out of memory
[ERROR 2022-11-30 21:00:25,847] worker.py: 400  Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): �[36mray::RayTrainWorker._RayTrainWorker__execute()�[39m (pid=9607, ip=172.31.5.55, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x7f9cca7c39a0>)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 31, in __execute
    raise skipped from exception_cause(skipped)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 28, in __execute
    return func(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/torch/config.py", line 118, in _shutdown_torch
    torch.cuda.empty_cache()
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/torch/cuda/memory.py", line 114, in empty_cache
    torch._C._cuda_emptyCache()
RuntimeError: CUDA error: out of memory
[ERROR 2022-11-30 21:00:25,859] worker.py: 400  Unhandled error (suppress with 'RAY_IGNORE_UNHANDLED_ERRORS=1'): �[36mray::RayTrainWorker._RayTrainWorker__execute()�[39m (pid=9605, ip=172.31.5.55, repr=<ray.train._internal.worker_group.RayTrainWorker object at 0x7f1191c979d0>)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 31, in __execute
    raise skipped from exception_cause(skipped)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/_internal/worker_group.py", line 28, in __execute
    return func(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/ray/train/torch/config.py", line 118, in _shutdown_torch
    torch.cuda.empty_cache()
  File "/home/ubuntu/miniconda3/envs/seeweed/lib/python3.8/site-packages/torch/cuda/memory.py", line 114, in empty_cache
    torch._C._cuda_emptyCache()
RuntimeError: CUDA error: out of memory

As mentioned in https://discuss.pytorch.org/t/out-of-memory-when-i-use-torch-cuda-empty-cache/57898/2, we need to empty the cache on the specific GPU device for the worker.

Versions / Dependencies

master

Reproduction script

Any TorchTrainer run

Issue Severity

Medium: It is a significant difficulty but I can work around it.