NCCL 버전 업그레이드 하다가 발생한 문제

Multi GPU를 사용하는 학습 코드를 돌리던 와중에 잘 돌아가다가 갑자기 특정 epoch에서 코드가 멈추는 현상이 발생했다.(GPU에도 메모리가 멈췄다. watch -n 0.1 nvidia-smi로 보는데, 계속 멈춰있다) 그러다가 갑자기 혼자 중단하고서 하는말이

First PG on this rank that detected no heartbeat of its watchdog. [rank1]:[E814 01:43:33.703972694 ProcessGroupNCCL.cpp:1413] [PG 0 (default_pg) Rank 1] Heartbeat monitor timed out! Process will be terminated after dumping debug info. workMetaList_.size()=1 964it [22:51, 1.42s/it][rank1]:[F814 01:53:33.705017387 ProcessGroupNCCL.cpp:1224] [PG 0 (default_pg) Rank 1] [PG 0 (default_pg) Rank 1] ProcessGroupNCCL's watchdog got stuck for 600 seconds without making progress in monitoring enqueued collectives. This typically indicates a NCCL/CUDA API hang blocking the watchdog, and could be triggered by another thread holding the GIL inside a CUDA api, or other deadlock-prone behaviors.If you suspect the watchdog is not actually stuck and a longer timeout would help, you can either increase the timeout (TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC) to a larger value or disable the heartbeat monitor (TORCH_NCCL_ENABLE_MONITORING=0).If either of aforementioned helps, feel free to file an issue to PyTorch about the short timeout or false positive abort; otherwise, please attempt to debug the hang.

한숨이 나왔다. 일단 저기서 추천하는 heartbeat timeout sec나 torch nccl enable monitoring 설정의 경우 문제를 완전히 해결해주지는 않는다. 전자의 경우는 시간을 조금더 늘려(갑자기 멈추는 시간)서 그 시간이 지날때까지 에러가 발생하지 않도록 기다려주는거고, 두 번째는 잘 모르겠다. 

찾아보니까 서버 GPU의 CUDA 버전이랑 NCCL(Nvidia collective communications library) 버전이 안맞아서 생길 수 있는 에러라고 한다.(수정: 결국 이게 원인은 아니었음. 진짜 원인은 맨 마지막에 공개)  내 NCCL 버전은  2.20.5였고, 이 버전과 호환되는 CUDA 버전은 12.4까지였으나, 내 CUDA 버전은 12.5였다.

 

그래서 NCCL 버전을 CUDA 12.5와 호환되도록 2.22.3으로 업그레이드 하였다. 업그레이드를 하는 방법은 여기에서 하면 된다. 일단 다운로드를 받고, 서버를 재부팅 했는데, 학습 코드 스크립트에 "export NCCL_DEBUG=INFO"를 통해 NCCL 버전을 확인한 결과, 업데이트가 반영되지 않은 것을 확인했다. (ㅠㅠ) 

혼자 삽질을 엄청하다가 정말 간단하게 해결할 수 있다는걸 깨닫고 나 같은 사람이 있을까봐 정리합니다.

일단 나는 최신 버전을 깔면 기존 버전은 알아서 지워진다고 생각했다. 그런데 이 생각이 이 문제의 원인이었다.

서버 자체가 연구실 사람들이 공통으로 쓰는 서버여서 그런지는 모르겠지만 NCCL 관련 파일이 엄청 많았고, 내가 방금 설치한 NCCL 버전은 어디에 있는지 감조차 잡을 수 없었다.

일단 팩트만 말하자면 NCCL을 새로 다운로드 받았다면 그 경로는 " /usr/lib/x86_64-linux-gnu/ " 여기이다. 일단, 이것이 잘 들어가 있는지  확인을 해야한다.

dpkg -L libnccl2
dpkg -L libnccl-dev

이 코드를 차례로 실행하고 나서, libnccl.so 나 libnccl.so.2가 " /usr/lib/x86_64-linux-gnu/ " 이 경로에 있는지 확인해야 한다. (나의 경우엔 잘 들어가 있었다) 

자 그럼 내 코드는 왜 자꾸 이전 버전의 NCCL을 가져오는걸까? 그래서 내 코드가 어떤 NCCL을 가져오는지 확인했다.

(JE) DAHS2@user:~/Timellm/Replicate_for_P$ LD_DEBUG=libs python run_EHRTimeLLM_v2(<-코드 파일 이름).py 2>&1 | grep libnccl
    125260:     find library=libnccl.so.2 [0]; searching
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cublas/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cuda_cupti/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cuda_nvrtc/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cuda_runtime/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cudnn/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cufft/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/curand/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cusolver/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/cusparse/lib/libnccl.so.2
    125260:       trying file=/home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2
    125260:     calling init: /home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2
    125260:     calling fini: /home/DAHS2/anaconda3/envs/JE/lib/python3.9/site-packages/torch/lib/../../nvidia/nccl/lib/libnccl.so.2 [0]

그 결과 마지막에 calling fini 에 있는 경로에서 가져오는걸 확인할 수 있다. 대체 왜 내 가상환경에 저게 들어가 있었는지는 모르겠지만,

저 파일들을 지우거나, 이름을 바꿔주어서 저 경로에서 가져오는걸 방지하면 문제가 해결된다

mv /home/DAHS2/anaconda3/envs/JE/bin/../lib/libnccl.so.2 /home/DAHS2/anaconda3/envs/JE/bin/../lib/libnccl.so.2.backup

 지우는건 쫄려서 이름을 바꾸었따.

이렇게 하고 나서 다시 학습 코드를 돌려보니, 아래와같이NCCL version이 최신 버전으로 나오는걸 확인했다.  

 

에러의 진짜 원인

NCCL과 CUDA 버전이 서로 호환 가능하도록 맞춰주었음에도, 특정 epoch에서 학습이 멈추는 현상은 똑같았다. 그래서 학습 과정 Log 기록을 유심히 관찰했고, 원인이 될 수 있는 부분을 발견했다.

참고로 로그에서 "you are going to heaven" 은 best validation loss가 현재 epoch의 validation loss보다 클 경우 모델을 저장한다는 의미에서 나타내는 print 문이다. GPU는 2개를 multi로 작동하고 있었다. 일단 epoch 7에서 you are going to heaven이 2번 시간 간격을 두고 등장했고, 모델이 저장됬다.

그러나 epoch 8에서는 early stoping count가 나왔는데, 바로 뒤에 you are going to heaven이 나왔다. 이 둘은 서로 베타적인 현상이고, 같은 시점에 일어나면 안된다.

...
````````````accelerator.print("Epoch: {} cost time: {}".format(epoch + 1, time.time() - epoch_time))
            train_loss = np.average(train_loss)
            vali_loss = vali(args, accelerator, model, vali_data, vali_loader, criterion)
            # test_loss = vali(args, accelerator, model, test_data, test_loader, criterion)
            accelerator.print(
                "Epoch: {0} | Train Loss: {1:.7f} Vali Loss: {2:.7f} ".format(
                    epoch + 1, train_loss, vali_loss))
            
            if vali_loss < best_loss:
                print('you are going to heaven')

                accelerator.wait_for_everyone() 
                if accelerator.is_local_main_process: 
                    model_to_save = accelerator.unwrap_model(model)
                    save_path = os.path.join(checkpoint_dir, f"Best_model_epoch_{epoch + 1}_loss-{vali_loss}.pt")
                    torch.save({
                        'epoch': epoch + 1,
                        'model_state_dict': model_to_save.state_dict(),
                        'optimizer_state_dict': model_optim.state_dict(),
                        'loss': train_loss,
                    }, save_path)
                    best_loss = vali_loss
                    accelerator.print(f"Model saved to {save_path}")
            early_stopping(vali_loss, model, path)
            gc.collect()
            torch.cuda.empty_cache()
            if early_stopping.early_stop:
                accelerator.print("Early stopping")
                break
                
...

코드를 잘못짰는가? 그건 또 아니다. validation loss를 계산하고, 그 후에 best loss와 비교하고, 그 다음 if 조건문으로 early stopping인지 확인한다.

내 생각엔 multi gpu가 어떻게 동작하는지는 모르겠지만, 두 GPU가 일관되지 않은 연산을 하고 있다고 생각했다.

그래서 vali loss와 best loss를 비교하는 부분을 아예 제거하고, early stopping class에서 그 부분을 넣었다. 사실 뭐가 문제였는지 정확하게 밝히지는 못했다.

class EarlyStopping:
    def __init__(self, accelerator=None, patience=7, verbose=False, delta=0, save_mode=True):
        self.accelerator = accelerator
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.best_score = None
        self.early_stop = False
        self.val_loss_min = np.Inf
        self.delta = delta
        self.save_mode = save_mode

    def __call__(self, val_loss, model, path):
        score = -val_loss
        if self.best_score is None:
            self.best_score = score
            if self.save_mode:
                self.save_checkpoint(val_loss, model, path)
        elif score < self.best_score + self.delta:
            self.counter += 1
            if self.accelerator is None:
                print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            else:
                self.accelerator.print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
            if self.counter >= self.patience:
                self.early_stop = True
        else:
            self.best_score = score
            if self.save_mode:
                self.save_checkpoint(val_loss, model, path)
            self.counter = 0

    def save_checkpoint(self, val_loss, model, path):
        if self.verbose:
            if self.accelerator is not None:
                self.accelerator.print(
                    f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')
            else:
                print(
                    f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')

        if self.accelerator is not None:
            model = self.accelerator.unwrap_model(model)
            torch.save(model.state_dict(), path + '/' + 'checkpoint')
        else:
            torch.save(model.state_dict(), path + '/' + 'checkpoint')
        self.val_loss_min = val_loss

Multi GPU를 활용하여 코드를 작성하는 것이 이렇게 어려울 줄 몰랐다. 아무튼 valid loss와 best loss를 비교하고, 모델을 저장하는 코드를 main code에서 완전히 제거하고, early stop 판단 시에 하도록 코드를 수정했더니, 이젠 에러가 나지 않았다.