PyTorch 2.14에서 학습 그래프에도 Inductor 지역성 재배치를 시험하는 방법 | DAKER 커뮤니티
PyTorch 2.14에서는 Inductor의 reorder_for_locality를 학습 그래프에서도 시험할 수 있습니다. 추론에서만 보던 지역성 재배치를 학습 컴파일 그래프에 적용해 볼 수 있다는 점이 이번 변화의 핵심입니다.
특히 torch.compile로 학습 루프를 묶어 쓰는 경우라면, 캐시 지역성이 아쉬운 구간이 실제로 병목인지 비교해 보기 좋습니다. 기본값은 꺼져 있으므로 기존 동작을 바꾸지 않은 상태에서, 같은 모델과 같은 배치 조건으로 켜기 전후의 스텝 시간과 메모리를 비교하면 됩니다.
PyTorch 2.14에서 reorder_for_locality_in_training를 켜면, 추론에만 적용되던 지역성 재배치를 학습 그래프에서도 옵트인으로 시험할 수 있습니다.
무엇이 달라졌나
이번에 확인할 설정은 reorder_for_locality_in_training입니다. 이 스위치를 켜면 Inductor의 지역성 재배치를 학습 컴파일 그래프에서도 시험할 수 있습니다. 다만 기본값은 off이므로, 성능 개선이 필요한 경우에만 따로 켜서 비교하는 방식이 적절합니다.
이 설정은 기본 동작을 강제로 바꾸지 않기 때문에 실험 조건을 맞추기 쉽습니다. 같은 입력, 같은 배치, 같은 컴파일 조건에서 on/off만 바꿔 보면 됩니다. 참가자 노트에는 PyTorch 버전, 컴파일 모드, 설정 on/off, 스텝 시간을 한 줄로 남겨 두는 것이 좋습니다.
재현해 볼 실험 조건
가장 간단한 방법은 작은 학습 스텝을 torch.compile로 감싼 뒤, 기준 실행과 비교 실행을 나란히 보는 것입니다.
- 작은 학습 스텝을
torch.compile로 감싼 뒤, 기본 설정으로 몇 스텝을 돌려 기준 시간을 적습니다. - Inductor 설정에서
reorder_for_locality_in_training만 켠 뒤 같은 입력을 다시 돌립니다. - 두 실행의 스텝 시간·예약 메모리를 한 표로 비교하고, 효과가 없으면 설정을 다시 끄면 됩니다. 추론 전용 지역성 패스와는 섞지 않는 편이 좋습니다.
simple_overlap은 집단통신과 연산을 겹치는 기본 재배치이고, 이번 글에서 다루는 것은 학습 그래프용 지역성 재배치 옵트인입니다. 실험이 끝나면 버전, 모델 크기, on/off 결과를 함께 남겨 두면 다음 사람이 같은 조건으로 바로 따라가기 좋습니다.
실무에서 바로 볼 포인트
이 설정은 학습 컴파일에서 캐시 지역성이 병목일 때만 따로 켜서 보는 것이 좋습니다. 기본값이 off이므로, 켠 뒤에는 반드시 끄기 전과 같은 조건으로 비교해야 의미 있는 결과를 얻을 수 있습니다.
또한 통신-연산 겹침 실험과 한 번에 섞지 않는 편이 좋습니다. 이번에는 지역성 스위치 하나만 바꿔 보고, 효과가 있는지 먼저 확인하면 됩니다.
이번 실험에서는 다른 최적화와 섞지 말고 reorder_for_locality_in_training on/off만 비교하는 편이 결과를 해석하기 쉽습니다.
관련 DAKER 학습
참고 자료
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/community
여러분은 torch.compile 학습 루프에서 이런 on/off 비교를 할 때 어떤 지표를 가장 먼저 보시는지 궁금합니다.