PyTorch 2.14의 torch.while_loop, CUDA 그래프로 가변 반복까지 한 번에 캡처하는 방법 | DAKER 커뮤니티

while_loop 그래프

한 줄 답: PyTorch 2.14부터 torch.while_loop는 CUDA while 조건부 노드로 CUDA 그래프에 캡처될 수 있습니다. 가변 trip count 구간을 곧바로 그래프 밖으로 밀지 않아도 되지만, 고정 최대 trip count와 텐서-only carried 입력 제약은 그대로입니다. 캡처 성공·동기화·스텝 시간만 같은 입력으로 비교하세요.

DAKER PyTorch 학습 독자를 위한 점검 노트입니다. 속도 배율은 워크로드·GPU마다 다르므로 숫자를 단정하지 않습니다.

무엇이 바뀌었다고 보나

이전에는 런타임에 정해지는 반복 횟수 때문에 캡처 경계를 쪼개는 경우가 많았습니다. 이제는 torch.while_loop 자체를 그래프 안에 넣는 선택지가 생겼습니다. 조건은 부모 스트림에서 평가되고 본문 끝에서 다시 검사됩니다.

가변 반복 때문에 그래프를 포기하던 구간을 다시 캡처 후보로 올릴 수 있습니다.

재현해 볼 실험 조건

  1. 고정 상한 trip count를 둔 torch.while_loop로 가변 길이 축소를 만듭니다.
  2. 같은 루프를 CUDA 그래프 캡처·재생으로 돌립니다.
  3. 캡처 실패 지점·디바이스-호스트 동기화·스텝 시간을 기록합니다.
  4. while_loop 없이 풀린 버전과 측정 축을 섞지 않습니다.

실무에서 바로 볼 포인트

성능 숫자보다 캡처 가능성 변화가 먼저입니다. 최대 trip count·텐서-only carried 제약을 노트에 적고, 빌드·CUDA 드라이버를 고정하세요.

측정표 템플릿

오늘 바로 할 측정

가변 길이 인덱스 루프 하나를 골라 캡처 전후만 비교합니다. Inductor simple_overlap·compile-on-one-rank와 축을 섞지 마세요.

마치며

torch.while_loop CUDA 그래프는 “켜면 무조건 이득”이 아니라 “가변 루프를 캡처 후보로 되돌리는 도구”입니다. 공식 릴리스 노트·PR을 기준으로 재현하세요.

관련 DAKER 글과 주제를 나누는 법

CUDA Graph 다중 풀·TunableOp GEMM은 인접 주제지만 측정 축이 다릅니다. while_loop 노트에는 max trip·캡처 성공·동기화만 남깁니다.

실패를 문서로 남기는 이유

캡처가 깨진 지점과 제약을 한 줄로 남기면, 다음 사람이 같은 루프를 다시 그래프 밖으로 밀어내지 않습니다.

고정 최대 trip count를 먼저 적는 이유

가변 반복이 허용돼도 캡처 경로에는 고정 상한이 필요합니다. 상한을 실험 노트 첫 줄에 적지 않으면, 나중에 “캡처 실패”만 보고 루프 전체를 그래프 밖으로 밀어내기 쉽습니다. 상한·실제 trip·실패 여부를 한 표에 나란히 둡니다.

carried 입력은 텐서-only

파이썬 객체나 호스트 동기화가 본문에 섞이면 캡처 경계가 다시 깨질 수 있습니다. while_loop 본문과 조건식에 들어간 입력이 텐서인지만 먼저 점검하세요. 호스트에서 길이를 읽어 오는 코드는 캡처 밖 전처리로 분리해 비교합니다.

출처