[PyTorch] Inductor mixed-order reduction 최적화

PR 링크: pytorch/pytorch#176495 상태: Merged | 변경: +133 / -1

들어가며

PyTorch Inductor의 mix-order-reduction은 여러 reduction 연산을 하나의 Triton 커널로 융합하는 최적화다. 그런데 multi-stage 파이프라이닝을 사용할 때, 큰 reduction 차원에서 GPU shared memory를 초과하여 컴파일이 실패하는 문제가 반복적으로 보고되었다. 이 PR은 multi-stage를 기본적으로 비활성화(NUM_STAGES=1)하여 안정성을 확보한다.

핵심 코드 분석

설정 플래그 추가

# torch/_inductor/config.py

# Don't allow multi-stages by default to avoid out of shared memory
mix_order_reduction_allow_multi_stages = (
    os.environ.get("TORCHINDUCTOR_MIX_ORDER_REDUCTION_ALLOW_MULTI_STAGES") == "1"
)

환경 변수로 opt-in 할 수 있는 설정을 추가했다. 기본값은 False로, 안전한 single-stage만 사용한다.

Heuristic 분기

Before:

MAX_NUM_STAGES = 2 if rnumel_hint > 8192 else 3
c.kwargs["NUM_STAGES"] = min(max(num_iters // 4, 1), MAX_NUM_STAGES)

After:

if inductor_meta.get("mix_order_reduction_allow_multi_stages", True):
    MAX_NUM_STAGES = 2 if rnumel_hint > 8192 else 3
else:
    MAX_NUM_STAGES = 1
c.kwargs["NUM_STAGES"] = min(max(num_iters // 4, 1), MAX_NUM_STAGES)

mix_order_reduction_allow_multi_stages가 False이면 MAX_NUM_STAGES를 1로 고정한다. 이전에는 reduction 차원(rnumel)이 8192 이하이면 최대 3-stage까지 허용했는데, 이것이 Eagle3 같은 대형 모델의 backward에서 shared memory 초과를 유발했다.

메타데이터 전달

# torch/_inductor/codegen/triton.py
def inductor_meta_common(cls):
    return {
        # ...
        "mix_order_reduction_allow_multi_stages":
            config.triton.mix_order_reduction_allow_multi_stages,
    }

Inductor codegen에서 runtime heuristic까지 설정값을 전달하는 경로를 추가했다.

왜 이게 좋은가

Multi-stage는 Triton 커널의 파이프라이닝 단계 수를 의미한다. 각 stage마다 shared memory에 중간 결과를 보관하므로, stage 수 x 데이터 크기가 GPU의 shared memory 한계(보통 48~164KB)를 넘으면 커널 컴파일이 실패한다. Eagle3 모델처럼 hidden size가 크고 RMSNorm + attention이 융합되는 경우에 이 문제가 발생했다.

기본값을 NUM_STAGES=1로 변경하면:

• Shared memory 초과로 인한 컴파일 실패가 사라진다
• 성능이 아주 약간 감소할 수 있지만, 안정성이 우선이다
• 고급 사용자는 환경 변수로 multi-stage를 다시 활성화할 수 있다

정리

• Mix-order-reduction의 NUM_STAGES 기본값을 1로 변경했다
• TORCHINDUCTOR_MIX_ORDER_REDUCTION_ALLOW_MULTI_STAGES=1로 기존 동작을 복원할 수 있다
• Eagle3 등 대형 모델의 backward에서 shared memory 초과 문제가 해결된다
• 36GB GPU 메모리를 사용하는 대규모 테스트 케이스가 추가되었다

참고 자료

• PyTorch Inductor 문서 -- Inductor 컴파일러 개요
• 이슈 #175250 -- 원본 shared memory 초과 버그 리포트

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.