PyTorch 2.14에서 ProcessGroup 생성 뒤 타임아웃을 바꾸는 방법 | DAKER 커뮤니티
분산 학습에서는 타임아웃이 너무 길어도, 너무 짧아도 문제가 됩니다. 한 랭크가 멈췄을 때 오래 기다리게 되기도 하고, 체크포인트 로드처럼 원래 시간이 걸리는 구간에서는 정상 작업이 중간에 끊기기도 합니다.
PyTorch 2.14의 torch.distributed.set_timeout은 이런 상황을 조금 더 유연하게 다룰 수 있게 합니다. ProcessGroup을 만든 뒤에도 집단통신 타임아웃을 바꿀 수 있어서, 느린 구간에서는 넉넉하게 두고 끝난 뒤에는 다시 짧게 돌리는 식으로 운영하면 됩니다.
PyTorch 2.14의 torch.distributed.set_timeout은 ProcessGroup을 만든 뒤에도 집단통신 타임아웃을 바꿉니다.
왜 초기화 뒤 타임아웃 변경이 중요한가
기본 타임아웃이 길면 장애 상황에서도 오래 대기하게 됩니다. 반대로 너무 짧으면 체크포인트 로드처럼 시간이 걸리는 구간에서 잡이 실패할 수 있습니다. 그래서 초기화 시점에 한 번 정한 값을 끝까지 유지하기보다, 구간에 따라 조정하는 편이 실무에 더 잘 맞습니다.
예전에는 _set_pg_timeout을 쓰던 자리를 이제 공개 API인 torch.distributed.set_timeout이 대신합니다. 이번 글의 초점도 여기에 있습니다. 그룹을 만든 뒤에 타임아웃을 다시 설정했을 때, 그 값이 이후 콜렉티브에 반영되는지를 확인하는 것입니다.
재현해 볼 실험 조건
실험은 작은 월드에서 시작하면 됩니다. 먼저 ProcessGroup을 열고 초기화 직후 타임아웃을 기록합니다. 그다음 torch.distributed.set_timeout으로 값을 늘린 뒤, 큰 상태 로드를 흉내 낸 sleep이나 실제 체크포인트 로드처럼 느린 구간을 한 번 둡니다. 이때 타임아웃 때문에 작업이 죽지 않는지 보면 됩니다.
이후에는 같은 API로 값을 다시 짧게 줄이고, 의도적으로 느린 콜렉티브가 더 빨리 실패하는지 비교하면 됩니다. 가능하면 연산별 타임아웃도 한 줄로 적어 두는 것이 좋습니다.
재현에는 행이 줄었다는 기록보다 초기화 뒤에 타임아웃을 바꿨는지가 더 도움이 됩니다.
참가자 노트에는 초기화 직후 값, 바꾼 값, 콜렉티브 종류를 남기면 됩니다. 실험이 끝나면 사용한 백엔드, 월드 사이즈, 타임아웃 초 값을 한 줄로 정리해 두면 다음 참가자가 같은 조건을 바로 따라오기 좋습니다.
이번 글에서 다루는 범위
이번 글은 그룹을 만든 뒤 타임아웃을 다시 설정하는 조건에만 초점을 둡니다. Flight Recorder 훅이나 nccl2 백엔드는 다른 글에서 다룹니다.
실무에서 바로 볼 포인트
타임아웃은 초기화 때만 정하는 값으로 두지 않는 것이 좋습니다. 느린 로드 구간과 일반 스텝 구간을 나눠 다르게 적용하면 운영이 수월해집니다. 또 학습 노트에는 변경 전후의 초 값과 콜렉티브 이름을 함께 남기면 재현과 비교에 도움이 됩니다. 비공개 _set_pg_timeout 대신 공개 API인 torch.distributed.set_timeout을 쓰는 것도 중요합니다.
느린 로드 구간에서는 타임아웃을 늘리고, 끝난 뒤에는 다시 짧게 둘 수 있습니다.
관련 DAKER 학습
참고 자료
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/community/post-mu5x4nfs-614089f7
https://daker.ai/community
실제로 운영하는 분들은 어떤 구간에서 타임아웃을 늘리거나 다시 줄이는 편인지 궁금합니다.