PyTorch 2.14부터 torch.while_loop는 CUDA while 조건부 노드로 CUDA 그래프 캡처가 가능합니다. 반복 횟수가 데이터에 따라 바뀌어도, 예전처럼 곧바로 graph break로 떨어지지 않는 경로가 열렸습니다.

자세한 내용과 근거는 DAKER 원문에서 확인하실 수 있습니다.

핵심 요약

• 동적 길이 루프를 compile에 넣다 보면, trip count가 입력에 묶여 그래프가 끊기는 장면을 자주 봅니다. 디코드처럼 길이가 샘플마다 다른 작업을 한 그래프에 남기고 싶을 때 특히 답답합니다. 오늘은 루프 본문이 매 스텝 같은 연산인지부터 확인하십시오. 본문이 안정적이면 while_loop를 그래프 안에 남길

• trip count가 입력 텐서에 의존하는 torch.while_loop를 준비합니다.

• CUDA에서 동일 입력을 고정 max trip 설정과 함께 그래프 캡처 경로로 돌립니다.

원문 보기: https://daker.ai/community/pytorch-2-14-torch-while-loop-cuda-graph-capture

DAKER: https://daker.ai/community/pytorch-2-14-torch-while-loop-cuda-graph-capture