torch.compiler.config.compile_on_one_rank로 한 랭크만 컴파일하면 다른 랭크는 같은 아티팩트를 로드합니다 | DAKER 커뮤니티
한 줄 요약
PyTorch 2.14의 torch.compiler.config.compile_on_one_rank는 분산 작업에서 모든 랭크가 같은 모델을 따로 컴파일하던 비용을, 한 랭크가 만든 아티팩트를 다른 랭크가 로드하는 경로로 바꿉니다. 옵트인이며, 프로그램당 가속기 한 대를 가정합니다.
월드 사이즈가 커질수록 “컴파일만 N번”이 시작 전 대기 시간을 늘니다. 오늘은 작은 멀티 랭크 잡에서 플래그를 켠 뒤, 한 랭크만 아티팩트를 쓰고 나머지가 같은 소스를 로드하는지부터 확인하십시오. 참가자 노트에는 플래그 이름, 디바이스 개수, 아티팩트 경로를 같이 남깁니다. “컴파일이 빨라졌다”보다 “누가 만들고 누가 로드했는지”를 적는 것이 재현에 도움이 됩니다.
재현해 볼 실험 조건
- 멀티 랭크 잡에서 모든 랭크가 독립적으로
torch.compile경로를 도는 경우와,torch.compiler.config.compile_on_one_rank를 켠 경우를 나란히 둡니다. - 한 랭크가 아티팩트를 생성·저장하고, 다른 랭크는 같은 아티팩트를 로드하는지 로그로 확인합니다. 코드젠이 디바이스에 묶이지 않는지(예: 한 디바이스에서 만든 커널을 다른 인덱스의 같은 종류 디바이스에서 로드)도 같이 봅니다.
- 그래프가 두 번째 가속기를 건드리면 거절되는지, 프로그램당 가속기 하나 가정이 깨지는 케이스를 짧게 적어 둡니다.
make_fx가 추적 랭크의 디바이스를 팩토리·캐스트에 박지 않는 점, Inductor·런처가 로드 시점에 디바이스를 고르는 점, DeviceMesh.get_group()이 메시에서 그룹을 가져오는 점은 같은 플래그 계열의 장치 중립 코드젠을 떠받칩니다. 이번 글에서는 “한 번 만들고 여러 랭크가 로드” 조건에만 초점을 둡니다. 실험이 끝나면 사용한 PyTorch 버전·월드 사이즈·플래그 값을 한 줄로 고정해 두면, 다음 참가자가 같은 조건을 바로 따라올 수 있습니다.
실무에서 바로 볼 포인트
- 플래그는 옵트인입니다. 기본값을 바꾸지 않은 채 켠 뒤에만 아티팩트 공유 경로를 기대하십시오.
- 프로그램당 가속기 하나 가정이 깨지면(한 프로세스가 두 디바이스를 동시에 씀) 이 모드가 맞지 않을 수 있습니다.
- 학습 노트에 “독립 컴파일 vs 단일 아티팩트 로드” 결과와 디바이스 인덱스를 남깁니다.
관련 DAKER 학습
설명용 생성 이미지입니다. 동작은 로컬 멀티 랭크 컴파일·아티팩트 로드 결과로 확인하십시오.