[딥러닝] MIG로 분할된 GPU 2개를 이용해 모델 학습 시 런타임 에러

문제 상황

학습 시작 시 아래와 같은 에러 발생

RuntimeError: device >= 0 && device < num_gpus INTERNAL ASSERT FAILED at "../aten/src/ATen/cuda/CUDAContext.cpp":50, please report a bug to PyTorch.

 

환경

- OS: Ubuntu 20.04.6 LTS

- Python: v3.8.10

- Torch: v2.0.1

- CUDA: v11.7

- GPU: A100 * 2 (MIG 적용)

nvidia-smi -L

 

해결

- Single GPU 학습

환경변수 CUDA_VISIBLE_DEVICES를 0으로 세팅해주니 해결됨

export CUDA_VISIBLE_DEVICES=0

 

- Multi GPU 학습

아직 CUDA 11과 12는 하나의 프로세스에 오직 하나의 MIG 인스턴스만 사용할 수 있으므로 불가능...

참고로 CUDA 11부터 A100 MIG, CUDA 12부터 H100 MIG 사용이 가능함

 

---

 

참고 1: https://github.com/pytorch/pytorch/issues/38616

RuntimeError: num_gpus <= 16 INTERNAL ASSERT FAILED · Issue #38616 · pytorch/pytorch

참고 2: https://docs.nvidia.com/datacenter/tesla/mig-user-guide/index.html#cuda-visible-devices

NVIDIA Multi-Instance GPU User Guide :: NVIDIA Tesla Documentation

참고 3: https://stackoverflow.com/questions/73175008

How to make multiple GPUs visible with os.environ["CUDA_VISIBLE_DEVICES"] using GPU_IDs