PyTorch 2.14에서 compile_on_one_rank로 분산 컴파일 중복을 줄이는 방법 | DAKER 커뮤니티

분산 학습이나 추론 작업에서는 모델 실행보다 먼저 컴파일 대기 시간이 체감되는 경우가 있습니다. 특히 월드 사이즈가 커질수록 같은 모델을 각 랭크가 따로 컴파일하는 비용이 시작 시간을 늘릴 수 있습니다.

PyTorch 2.14의 torch.compiler.config.compile_on_one_rank는 이런 상황에서 눈여겨볼 만한 옵션입니다. 한 랭크가 만든 아티팩트를 다른 랭크가 로드하는 경로를 열어, 같은 컴파일을 반복하는 부담을 줄이는 방식입니다.

PyTorch 2.14의 torch.compiler.config.compile_on_one_rank는 분산 작업에서 모든 랭크가 같은 모델을 따로 컴파일하던 비용을, 한 랭크가 만든 아티팩트를 다른 랭크가 로드하는 경로로 바꿉니다.

다만 이 기능은 기본 동작이 아니라 옵트인이며, 프로그램당 가속기 한 대를 가정합니다. 그래서 단순히 빨라졌는지만 보기보다, 어떤 랭크가 아티팩트를 만들고 다른 랭크가 무엇을 로드했는지까지 함께 확인하는 것이 재현에 도움이 됩니다.

compile_on_one_rank가 바꾸는 점

기본적으로 멀티 랭크 잡에서는 각 랭크가 독립적으로 torch.compile 경로를 돌 수 있습니다. 이 경우 월드 사이즈가 커질수록 사실상 같은 컴파일이 여러 번 반복됩니다.

torch.compiler.config.compile_on_one_rank를 켜면 기대하는 경로는 달라집니다. 한 랭크가 아티팩트를 생성하고 저장한 뒤, 다른 랭크는 같은 소스를 다시 컴파일하는 대신 그 아티팩트를 로드합니다.

중요한 것은 컴파일 시간이 줄었는지보다 누가 아티팩트를 만들고 누가 같은 결과물을 로드했는지 확인하는 일입니다.

참가자 노트나 실험 기록에는 플래그 이름, 디바이스 개수, 아티팩트 경로를 함께 남겨 두는 것이 좋습니다. 같은 조건을 다시 맞출 때 도움이 되기 때문입니다.

재현해 볼 실험 조건

이 기능을 확인할 때는 두 경로를 나란히 두고 비교하면 됩니다. 하나는 멀티 랭크 잡에서 모든 랭크가 독립적으로 torch.compile 경로를 도는 경우이고, 다른 하나는 torch.compiler.config.compile_on_one_rank를 켠 경우입니다.

이때 핵심은 로그 확인입니다. 한 랭크가 아티팩트를 생성·저장하고, 다른 랭크가 같은 아티팩트를 로드하는지 살펴보면 됩니다. 또한 코드젠이 디바이스에 묶이지 않는지도 함께 보는 것이 좋습니다. 예를 들어 한 디바이스에서 만든 커널을 다른 인덱스의 같은 종류 디바이스에서 로드하는 상황을 확인하는 식입니다.

아울러 그래프가 두 번째 가속기를 건드릴 때 거절되는지도 짧게 적어 두면 좋습니다. 이 기능은 프로그램당 가속기 하나를 가정하므로, 그 전제가 깨지는 경우에는 맞지 않을 수 있습니다.

장치 중립 코드젠과 이번 글의 범위

make_fx가 추적 랭크의 디바이스를 팩토리·캐스트에 박지 않는 점, Inductor·런처가 로드 시점에 디바이스를 고르는 점, DeviceMesh.get_group()이 메시에서 그룹을 가져오는 점은 같은 플래그 계열의 장치 중립 코드젠을 떠받칩니다.

다만 이번 글에서는 그 전체를 넓게 다루기보다, 한 번 만든 아티팩트를 여러 랭크가 로드하는 조건에만 초점을 둡니다. 실험이 끝나면 사용한 PyTorch 버전, 월드 사이즈, 플래그 값을 한 줄로 정리해 두면 다음 사람이 같은 조건을 바로 따라가기 좋습니다.

실무에서 바로 볼 포인트

이 기능은 옵트인입니다. 기본값을 바꾸지 않은 상태에서는 아티팩트 공유 경로를 기대하기 어렵습니다. 따라서 플래그를 켠 경우와 켜지 않은 경우를 분리해서 기록하는 것이 좋습니다.

또한 한 프로세스가 두 디바이스를 동시에 쓰는 식으로 프로그램당 가속기 하나 가정이 깨지면, 이 모드가 맞지 않을 수 있습니다. 이런 경우에는 결과만 적기보다 어떤 디바이스 인덱스를 사용했는지까지 남겨 두면 해석이 쉬워집니다.

학습 노트에는 독립 컴파일과 단일 아티팩트 로드 결과를 함께 적어 두는 편이 좋습니다.

참고 자료

https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/public/learning/materials/pytorch-advanced-2-embedding-self-attention-mini-i
https://daker.ai/community

여러분은 멀티 랭크 환경에서 컴파일 시간을 볼 때, 속도 자체보다 아티팩트 생성·로드 경로를 얼마나 함께 기록하고 계신가요?