PyTorch Inductor에서 반복문의 상태 변경 보존하기

함수가 돌려준 값은 맞는데 같은 함수에 같은 텐서를 다시 넣으면 결과가 틀릴 수 있을까요?

입력 텐서(tensor)의 값을 함수가 직접 바꾸는 경우입니다. 반환값을 올바르게 계산하더라도 원래 텐서에 변경을 남기지 못하면 다음 호출은 잘못된 값에서 시작합니다.

PyTorch의 이슈 #195327은 torch.while_loop를 컴파일할 때 나타나는 이런 오류를 다룹니다. 2026년 9월 4일 PyTorch에 반영된 PR #195393의 핵심은 반복문의 본문뿐 아니라 조건을 검사하면서 바꾸는 값도 컴파일러가 추적해야 한다는 것입니다.

1. 반환값과 원래 텐서는 따로 확인해야 한다

다음 코드는 텐서의 값을 절반으로 바꾸고, 그 결과를 복제해서 반환합니다.

import torch

def halve(state):
    state.mul_(0.5)
    return state.clone()

state = torch.tensor(1.0)
answer = halve(state)

mul_는 원래 텐서의 값을 직접 바꿉니다. 이런 동작을 제자리 변경(mutation)이라고 합니다. 밑줄이 없는 mul은 원래 텐서를 그대로 두고 곱셈 결과를 새 텐서로 반환합니다.

마지막의 clone()은 별도 저장 공간에 현재 값을 복제합니다. 따라서 함수가 반환한 answer와 호출자가 가진 state는 따로 확인할 수 있습니다. 위 코드에서는 둘 다 0.5여야 합니다.

컴파일한 함수가 answer만 0.5로 만들고 state를 1.0으로 남겨 둔다면 어떨까요? 반환값만 검사하는 테스트는 통과합니다. 하지만 같은 state로 다시 호출하면 기대한 0.25가 나오지 않습니다.

컴파일러는 함수가 돌려주는 값뿐 아니라 입력에 남기는 변경도 보존해야 합니다. 이슈의 반복문에서도 바로 이 문제가 생깁니다.

2. 반복문의 조건도 값을 바꿀 수 있다

보통 반복문에서 값을 바꾸는 곳은 본문이라고 생각하기 쉽습니다. 하지만 조건을 검사하는 함수도 값을 바꿀 수 있습니다.

torch.while_loop는 “계속 반복할까?”를 판단하는 cond와 “이번에 무엇을 할까?”를 실행하는 body를 따로 받습니다. 다음은 이슈의 핵심을 간추린 함수입니다. torch를 가져온 뒤 추론 모드인 torch.inference_mode()에서 호출하는 예제입니다.

def f(state):
    state.mul_(0.8)

    def cond(i, acc):
        state.mul_(0.5)
        return i < 2

    def body(i, acc):
        return i + 1, acc + state

    i, acc = torch.while_loop(
        cond,
        body,
        (
            torch.zeros((), dtype=torch.int64, device=state.device),
            torch.zeros_like(state),
        ),
    )
    return i, acc, state.clone()

처음에는 코드에서 세 군데만 보면 됩니다.

• 반복문에 들어가기 전에 state에 0.8을 곱합니다.
• 조건 함수 cond는 호출할 때마다 state를 절반으로 줄입니다.
• 본문 함수 body는 현재 state를 누적값 acc에 더합니다.

i는 본문을 실행한 횟수입니다. i가 2보다 작을 때만 본문을 실행하므로 본문은 두 번 실행합니다. acc는 그동안 더한 값을 담습니다.

여기서 cond는 인자 목록에 없는 state를 사용합니다. 안쪽 함수가 바깥 함수의 변수를 참조하는 클로저(closure)이기 때문입니다. 텐서를 복제해 사용하는 것이 아니라 같은 텐서를 참조하므로 state.mul_(0.5)는 바깥의 텐서 값을 바꿉니다.

두 번 반복해도 조건 검사는 세 번

state가 1.0에서 시작할 때 값을 따라가 보겠습니다.

두 번째 본문을 실행한 뒤에도 반복을 끝낼지 판단해야 합니다. cond를 한 번 더 호출하므로 그 안의 state.mul_(0.5)도 실행합니다. 조건의 결과가 거짓이라고 해서 이미 수행한 곱셈이 취소되지는 않습니다.

따라서 함수가 끝나면 누적값은 약 0.6, state는 0.1입니다. 소수 계산은 읽기 쉽게 반올림했습니다.

이 예제에서는 i, acc처럼 다음 반복에 넘기는 값을 반복 전달 입력(carried inputs)이라고 합니다. state는 여기에 포함되지 않고 바깥에서 참조하는 텐서입니다. 분석 대상 버전의 while_loop 제약 사항은 추론 중 이런 외부 텐서의 변경을 허용합니다. 반복 전달 입력 자체의 제자리 변경이나 외부 파이썬 사전·리스트의 변경은 허용하지 않습니다.

외부 텐서와 반복 전달 입력을 함께 변경하면 검사가 우회될 수 있다는 후속 이슈 #195966도 보고되어 있습니다. 이 글은 문서에서 허용하는 외부 텐서만 변경하는 경우로 범위를 한정합니다.

3. 첫 호출에서는 숨고 두 번째 호출에서 드러나는 오류

이슈의 실행 결과를 보면 일반 실행과 문제가 있는 Inductor 실행의 차이가 드러납니다. Inductor는 torch.compile에서 텐서 연산을 최적화하고 실행 코드를 만드는 컴파일러입니다.

첫 호출에서 Inductor가 반환한 누적값과 상태 복제본은 모두 맞습니다. 그런데 호출자가 가진 원래 텐서는 반복문 직전 값인 0.8로 남습니다. 반복문 안에서 계산한 0.1을 원래 텐서에 제대로 반영하지 못한 것입니다.

이 차이 때문에 두 번째 호출의 입력이 달라집니다. 정상 실행은 0.1에서 시작하지만 문제가 있는 실행은 0.8에서 시작합니다. 두 번째 누적값도 각각 0.06과 0.48로 갈립니다.

이 예제의 정상 실행은 호출할 때마다 상태를 이전 값의 10분의 1로 줄입니다. 1.0 → 0.1 → 0.01처럼 이어져야 하는데 첫 호출 뒤부터 이 흐름이 끊깁니다. 반환값 검사만으로는 놓치는 이유입니다.

4. 원인: 계산을 마치기 전에 원래 텐서에 복사한다

왜 내부 계산은 맞는데 원래 텐서에는 다른 값이 남을까요?

컴파일러는 제자리 변경을 분석하기 편하도록 원래 텐서와 별도로 계산한 값을 사용할 수 있습니다. 입력을 직접 바꾸는 연산을 새 값을 만드는 연산으로 바꾸는 과정을 함수형 변환(functionalization)이라고 합니다. 계산을 마치면 최종 값을 원래 텐서에 다시 복사해야 합니다.

이 사례의 필요한 순서를 값과 함께 그리면 다음과 같습니다. 실제 생성 코드가 아니라 동작을 설명하는 도식입니다.

정상적인 순서

1. 반복문 전 계산: 1.0 × 0.8 = 0.8
2. 반복문 안 계산: 0.8 → 0.4 → 0.2 → 0.1
3. 원래 텐서에 복사: 0.1

마지막 복사에는 copy_ 같은 연산을 사용합니다. clone()이 새 저장 공간을 만드는 것과 달리, copy_는 기존 텐서의 저장 공간에 값을 씁니다.

PR의 원인 분석에 따르면 문제가 있는 경로에서는 컴파일러가 조건 함수의 변경을 놓칩니다. 실행 순서를 정하는 스케줄러(scheduler)는 반복문이 이 값을 읽기만 한다고 판단합니다. 값을 바꾸지 않는다고 판단했으므로 복사를 반복문 뒤까지 미룰 이유도 없습니다.

그 결과 원래 텐서에 복사하는 연산을 반복문 앞의 계산과 합쳐 버립니다.

문제가 있는 순서

1. 반복문 전 계산: 1.0 × 0.8 = 0.8
   원래 텐서에 복사: 0.8
2. 반복문 안 계산: 0.8 → 0.4 → 0.2 → 0.1
   내부 계산 결과는 맞지만 원래 텐서는 여전히 0.8

컴파일러에 알려야 할 관계는 “반복문이 값을 바꾸므로, 그 값의 복사는 반복문이 끝난 뒤에 한다”입니다. 이런 실행 순서의 제약을 의존성(dependency)이라고 합니다. 곱셈 자체가 틀린 것이 아니라, 복사 시점을 정하는 데 필요한 정보가 빠진 문제입니다.

5. 수정: 조건과 본문이 바꾸는 텐서를 모두 기록한다

PR은 조건 함수의 변경을 찾아 그 사실을 실행 순서에 반영합니다.

조건 함수도 조사하기

기존 코드는 본문이 어떤 입력을 바꾸는지 조사합니다. 이 예제의 본문은 state를 읽기만 하므로 본문만 보면 state의 변경을 찾지 못합니다.

수정안은 조건 함수와 본문을 모두 조사하고 결과를 합칩니다.

조건 함수가 바꿀 수 있는 텐서: state
본문이 바꿀 수 있는 텐서:     없음
반복문 전체에서 기록할 텐서: state

어느 한쪽에서라도 바꾸는 텐서를 모두 모으는 합집합(union) 방식입니다. 두 함수가 같은 텐서를 바꿔도 목록에 한 번만 기록하면 됩니다.

찾은 변경을 실행 순서에 반영하기

변경을 찾은 뒤에는 스케줄러에도 그 사실을 알려야 합니다.

Inductor는 MutationOutput이라는 내부 기록으로 “이 연산은 기존 텐서의 저장 공간을 바꾼다”는 사실을 나타냅니다. 이름에 Output이 있지만 사용자에게 새 반환값을 하나 더 돌려준다는 뜻은 아닙니다. MutationOutput 구현도 새 저장 공간을 할당하지 않고 기존 공간의 변경을 등록합니다.

PR은 WhileLoop.create에서 조건과 본문 양쪽의 변경을 찾습니다. state 같은 외부 입력의 변경도 이 방식으로 등록합니다. 그러면 스케줄러는 원래 텐서에 복사하는 연산을 반복문 앞으로 옮기지 않도록 필요한 의존성을 표현할 수 있습니다.

6. 확인: 같은 텐서로 다시 호출해 보기

이 오류를 잡으려면 컴파일하지 않은 실행과 컴파일한 실행에서 다음 두 가지를 함께 비교해야 합니다.

• 함수가 반환한 값이 같은가?
• 호출이 끝난 뒤 원래 텐서의 값도 같은가?

이어서 텐서를 초기화하지 않고 같은 함수에 다시 넣어 봐야 합니다. 매번 새 텐서로 시작하면 앞선 호출이 잘못 남긴 상태를 다음 호출에서 사용하는지 확인하지 못합니다.

PR의 반복 호출 테스트는 같은 상태로 세 번 호출하면서 반환값과 원래 텐서를 매번 비교합니다.

본문을 한 번도 실행하지 않는 경우도 중요합니다. PR의 테스트처럼 조건 함수가 거짓인 텐서를 반환하면 본문은 실행하지 않아도 조건 함수는 한 번 실행합니다. 위 예제에서 조건을 그렇게 바꾸면 state는 1.0 × 0.8 × 0.5 = 0.4가 되어야 합니다. 반면 파이썬 상수 False를 반환할 때는 Dynamo가 조건 함수의 외부 텐서 변경을 누락하는 별도 이슈 #195970이 남아 있습니다.

맺음말

이슈의 핵심은 첫 반환값이 맞아도 다음 호출에 쓸 텐서의 값은 틀릴 수 있다는 점입니다. PR은 조건 함수의 변경까지 찾아 기록해, 반복문이 계산을 마친 뒤 원래 텐서에 결과를 복사하도록 합니다.

상태를 계속 갱신하는 코드를 검사할 때는 반환값과 함께 이번 호출이 다음 호출에 어떤 값을 남기는지도 확인해야 합니다.

이 글은 2026년 9월 7일자 포스트입니다. 분석 기준은 PyTorch에 반영된 fc660a81ddb088588f79920decaa9aabe7734cb7 커밋과 관련 논의입니다. PyTorch 재현 코드나 PR 테스트를 새로 실행하지는 않았습니다. 실행 결과는 이슈·PR의 보고를 인용하며 구현 설명은 해당 코드와 원인 분석을 바탕으로 합니다. AI의 도움으로 조사하고 작성한 글입니다.