그래프멈드 파이토치2 그래프 브레이크 자동 제거 기술

초록

GRAPHMEND는 파이토치2의 FX 그래프 트레이싱 과정에서 발생하는 동적 제어 흐름과 파이썬 부작용에 의한 그래프 브레이크를 고수준 코드 변환으로 사전에 제거한다. Jaseci 기반 AST 변환을 적용해 모델을 단일 연속 FX 그래프로 재구성함으로써 CPU‑GPU 동기화 비용을 크게 낮추고, 8개 Hugging Face 모델에서 최대 75%의 레이턴시 감소와 8% 수준의 처리량 향상을 달성한다.

상세 분석

본 논문은 파이토치2가 도입한 TorchDynamo‑TorchInductor JIT 파이프라인이 동적 제어 흐름(data‑dependent control flow)이나 파이썬 사이드 이펙트(I/O, print 등)를 만나면 FX 그래프를 중단하고 eager 모드로 전환하는 ‘그래프 브레이크’ 현상을 근본적으로 해결하고자 한다. 기존 접근법은 바이트코드 수준에서 지원되지 않는 연산을 탐지해 즉시 fallback하는 방식으로, 고수준 프로그램 구조(조건문, 루프, 부작용) 정보를 상실한 채 부분적인 최적화만 가능했다. GRAPHMEND는 Jaseci 컴파일러 프레임워크 위에 AST‑level 변환 패스를 삽입함으로써, 파이썬 소스 코드를 분석·재구성한다. 두 가지 핵심 변환은 (1) ‘Predicated Dynamic Control Flow’로, 데이터‑의존 조건문을 torch.where·torch.cond와 같은 텐서 연산으로 대체해 조건 평가 자체를 그래프 내부에서 수행하도록 만든다. 이는 조건문이 런타임 텐서 값에 의존하더라도 FakeTensor 기반의 정적 트레이싱이 가능하게 하여, 기존에 삽입되던 그래프 브레이크를 제거한다. (2) ‘Graph Epilogue Deferred Side Effects’는 print·logging 등 부작용을 함수 종료 시점으로 지연시켜, 트레이싱 단계에서는 순수 연산만 남게 만든다. 이때 부작용은 버퍼에 저장했다가 그래프 실행 후 한 번에 플러시한다. 변환 과정은 AST → UniIR → CFG 통합 분석을 통해 제어 흐름과 데이터 의존성을 정확히 파악하고, 변환 전후의 의미적 동등성을 보장한다. 실험에서는 8개의 Hugging Face 모델을 대상으로 그래프 브레이크 수를 6개 모델에서 0으로, 또 다른 모델에서는 5→2로 감소시켰으며, 남은 1개 모델은 현재 변환으로 해결 불가능한 복합 패턴을 포함한다. 성능 측면에서는 RTX 3090과 A40 GPU에서 콜드 스타트 레이턴시가 30‑75% 감소하고, 스테디‑스테이트 레이턴시가 2.5‑25% 낮아졌다. 또한 전체 파이프라인 처리량이 5‑8% 향상되었다. 이러한 결과는 고수준 코드 변환이 파이토치2 JIT의 한계를 보완하고, 개발자가 기존 코드를 그대로 유지하면서도 최적화된 GPU 실행을 얻을 수 있음을 입증한다. 논문은 또한 변환이 어려운 복합 제어 흐름(중첩 루프·조건·부작용 혼합)과 현재 Jaseci IR이 지원하지 않는 일부 파이썬 메타프로그래밍 패턴에 대한 향후 연구 방향을 제시한다.