IL-0125 BigScience · 2022-05-18ai-infra

Jean Zay 에서 진행한 176B 다국어 모델 학습에서 GPU 가 고장 나도 학습 작업이 끝나지 않고 멈춰 버려 SLURM 이 작업을 자동으로 재시작하지 못했다.

BigScience 는 Jean Zay 슈퍼컴퓨터에서 176B 파라미터 다국어 모델 tr11-176B-ml 을 학습했다. 학습 중 대략 일주일에 한 번 GPU 하나가 죽었다. 대부분은 SLURM 작업이 자동으로 재시작됐고, 체크포인트를 100 이터레이션마다 저장했기 때문에 잃는 학습은 최대 3시간이었다. 100 이터레이션을 학습하는 데 약 3시간이 걸렸다. 저장 간격은 처음에 더 길었고, 2022-03-21 GPU 고장으로 7.3시간을 잃은 뒤 200 이터레이션으로, 이어서 100 이터레이션으로 줄였다. 체크포인트 하나가 2.3TB 이고 저장에 40초가 걸려 너무 자주 저장할 수는 없었다.

문제는 프로세스가 죽었는데도 학습이 멈춘 채 끝나지 않는 경우였다. 이렇게 되면 SLURM 작업도 끝나지 않아서, 누군가 알아채지 않는 한 학습이 재시작되지 못했다. 누군가 하루 종일 학습을 지켜봐야 했고, 팀이 만든 킬 스위치도 가끔 동작하지 않았다. Jean Zay 에서는 sudo 권한이 없고, 한 사람이 띄운 SLURM 작업을 다른 사람이 멈출 수 없다. 그래서 팀은 특정 파일을 만들면 체크포인트를 저장한 뒤 학습을 정상 종료하는 킬 스위치를 따로 만들어 쓰고 있었다.

작성자가 이 문제를 PyTorch 이슈 76287 로 보고하자, PyTorch 쪽은 브로드캐스트 멈춤 같은 NCCL 문제가 생기면 강제로 죽도록 NCCL_ASYNC_ERROR_HANDLING=1 을 설정하고 srun 에 --wait=60 과 --kill-on-bad-exit=1 을 주라고 제안했다. --wait=60 은 첫 태스크가 끝나고 60초 뒤 나머지 태스크를 모두 끝내는 옵션이고, --kill-on-bad-exit=1 은 태스크 하나라도 0 이 아닌 종료 코드로 끝나면 스텝 전체를 끝내는 옵션이다. 노드 2개에서 시험해 보니 프로세스 하나를 죽이고 약 5분 뒤 SLURM 이 작업 전체를 끝냈다. 나중에 Jean Zay 의 SLURM 에는 이미 UnkillableStepTimeout=180, KillWait=60, KillOnBadExit=1 이 설정되어 있었다는 것을 알았고, 실제로 새로 더한 것은 NCCL_ASYNC_ERROR_HANDLING=1 하나였다.

심각도

–

영향 범위

The tr11-176B-ml training run; up to about 3 hours of training lost per GPU failure, and hung jobs stayed idle until someone noticed

영향 사용자

–

지역

–

감지까지

–

복구까지

–