RuntimeError: Cannot re-initialize CUDA in forked subprocess

실행 환경

Celery

$ celery -A utils.celery.celery_app worker --hostname=gen --queues=gen_queue --loglevel=info 

Celery Task

@celery_app.task(name="some_gen_func", queue="gen_queue")
def some_gen_func(
        model, ...
):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pipe = StableDiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16).to(device)
		...

에러 메시지

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method
[2024-09-29 18:03:04,845: ERROR/ForkPoolWorker-125] Task some_gen_func[...task_id] raised unexpected: RuntimeError("Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method")
Traceback (most recent call last):
  File "/home/user/.local/lib/python3.10/site-packages/celery/app/trace.py", line 453, in trace_task
    R = retval = fun(*args, **kwargs)
  File "/home/user/.local/lib/python3.10/site-packages/celery/app/trace.py", line 736, in __protected_call__
    return self.run(*args, **kwargs)
  File "/home/user/project/S11P21S001/ai/utils/tasks.py", line 31, in text_to_image_task
    torch.cuda.set_device(gpu_device)
  File "/home/user/.local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 350, in set_device
    torch._C._cuda_setDevice(device)
  File "/home/user/.local/lib/python3.10/site-packages/torch/cuda/__init__.py", line 235, in _lazy_init
    raise RuntimeError(
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

해결

• --pool=threads 로 셀러리 워커 실행

$ celery -A utils.celery.celery_app worker --hostname=gen --queues=gen_queue --loglevel=info --pool=threads

멀티스레딩으로 해결한 이유

1. PyTorch의 멀티프로세싱 방식(spawn)과 셀러리 멀티프로세싱(fork) 방식의 충돌로 인해 멀티프로세싱이 불가능

현재 우리 프로젝트에서 셀러리를 사용하는 목적은 무거운 AI 작업을 병렬로 실행하여 서버의 응답 지연을 최소화하고 자원을 효율적으로 활용하기 위함이었다. 따라서 처음 셀러리 구현을 기획할 때 막연하게 우리의 작업이 CPU-bound 작업이라 생각하고 Multi Process로 구현해야겠다고 생각했다.

또한 Multi Process로 구현 시의 이점도 있었다. --max-tasks-per-child 옵션을 주어 메모리 누수를 방지할 수 있기 때문이다. --max-tasks-per-child=1 옵션을 주어 셀러리를 시작하면 각 프로세스는 1개의 작업을 완료할 때마다 종료되고 새 프로세스를 자동으로 시작할 수 있다. 우리는 공용 GPU 서버를 사용하고 있었기 때문에 메모리 누수를 방지하는 것이 필수적이었다. 해당 옵션을 사용하면서 메모리에 대해 전혀 걱정하지 않고 있었다.

그러나 갑자기(분명 전날엔 똑같은 환경에서 실행됐을 때 잘 돌아갔는데 왜 도대체 왜) 해당 에러가 뜨면서 계획을 전부 뒤집어 엎어야 했다.

조사해보니 에러 메시지에 쓰여진 그대로, CUDA의 메모리 관리 방식은 fork 방식과 호환되지 않는다고 한다. fork는 부모 프로세스의 모든 메모리 상태를 그대로 복사하기 때문에 GPU 메모리(CUDA)와 관련된 자원들도 복사하려 시도한다. 그러나 GPU 메모리와 프로세스 메모리 간의 공유가 복잡하기 때문에 이를 복사하는 과정에서 다음과 같은 문제들이 발생할 수 있다.

• CUDA 초기화 문제: CUDA는 단 한 번만 초기화될 수 있다. 하지만 fork 방식에서는 부모 프로세스가 이미 CUDA를 초기화한 상태에서 자식 프로세스가 이를 복사하려 하기 때문에 CUDA를 재초기화하려는 시도가 이루어진다. 이것이 RuntimeError: Cannot re-initialize CUDA in forked subprocess와 같은 에러로 이어진다.
• 메모리 비일관성: fork 방식으로 GPU 자원을 사용할 때, 자식 프로세스가 부모 프로세스의 GPU 메모리 상태를 복사하기 때문에 비일관성이 발생할 수 있다. 이는 자식 프로세스가 부모와 독립적인 실행 환경을 가지지만, GPU 메모리를 공유하려고 하면서 충돌이 발생하는 것이다.
• 복잡한 메모리 공유 구조: GPU는 공유 메모리 방식이 복잡하여, fork 방식이 이를 잘 처리하지 못한다. 부모 프로세스와 자식 프로세스 간에 GPU 메모리를 명확하게 분리하여 관리할 수 있는 spawn 방식을 사용해야만 이러한 문제를 해결할 수 있다.

spawn 방식은 부모 프로세스의 메모리 상태를 복사하지 않고 완전히 새로운 프로세스를 생성하는 Multiprocessing 방식이다. 그러나 셀러리에서는 spawn 방식을 지원하지 않는다. 이러한 이유로 셀러리 워커를 멀티프로세스 환경에서 실행할 수 없게 되었다.

2. GPU-bound 작업은 파이썬의 Multi-threading 방식으로도 효과적인 처리가 가능

위에서 언급했던 것처럼 나는 우리의 Celery Task가 막연히 CPU-bound한 작업이라고 생각하고 있었다. 내 머릿속에서는 단순히 IO-bound vs CPU-bound 의 이분법적인 개념만을 염두에 두고있었기 때문이다.

그러나 우리의 작업은 GPU 서버에서 실행되며, 본질적으로 GPU-bound 작업이다. 우리의 Celery Task는 딥러닝 모델 StableDiffusionPipeline을 이용한 텐서 연산과 이미지 생성 등, 대부분의 연산을 GPU에서 수행한다.

앞서 말했듯이 멀티 프로세싱이 불가능해지면서 어쩔 수 없이 --pool=thread로 워커를 실행한 후 작업을 동시에 실행시켜봤는데, 내 예상과는 다르게 작업이 병렬적으로 잘 실행되었다. 병렬처리가 되지 않을 것이라고 생각했던 이유는 파이썬의 GIL 때문이었다.

Python에서는 GIL이라는 메커니즘이 있어서 멀티 스레드를 사용할 때도 한 번에 하나의 스레드만 Python 바이트코드를 실행할 수 있다. (자세한 내용은 Python의-GILGlobal-Interpreter-Lock이란 참조) 이는 CPU-bound 작업에 있어서 멀티 스레드의 성능을 제한하는 요인이 된다.

그러나 GPU-bound 작업에서는 상황이 달라진다. GPU 환경에서는 GIL의 제약이 크게 영향을 미치지 않기 때문이다!

우리의 Python 코드는 GPU에 작업을 전달할 뿐 연산은 대부분 GPU 메모리에서 수행된다. 또한 PyTorch와 같은 라이브러리는 GPU 작업을 비동기적으로 처리한다. 즉, 작업이 GPU로 전달되면 CPU에서는 해당 작업을 기다리지 않고 즉시 다른 작업을 수행할 수 있다. 이러한 이유로 우리의 작업들은 병렬적으로 처리될 수 있었던 것이다.

물론 가능하다면 멀티프로세싱의 이점이 훨씬 크다. 멀티프로세싱 시에는 각 프로세스가 독립적으로 GPU 메모리와 연산 리소스를 요청할 수 있기 때문에 더 효율적으로 GPU를 활용하고 병렬성을 보다 잘 관리할 수 있다.

그러나 현재 상황에서 멀티 프로세싱을 구현하는 것이 불가능하기 때문에 멀티 스레딩이 최선의 방식이 되었다. 다행히 구현해놓은 Celery를 활용할 수 없는 최악의 상황은 막았지만, 프로세스 재시작 등의 메모리 관리상의 이점을 사용할 수 없는 점은 앞으로 고민해야 할 숙제다.