GraphMend: Code Transformations for Fixing Graph Breaks in PyTorch 2

기술 요약: GraphMend

문제 정의

PyTorch 2는 TorchDynamo와 TorchInductor를 활용하여 모델 연산을 FX 그래프를 통해 최적화된 GPU 코드로 변환하는 JIT(Just-In-Time) 컴파일 파이프라인을 도입했습니다. 그러나 이 파이프라인은 지원되지 않는 Python 연산(예: 데이터 의존적 제어 흐름, 로깅과 같은 부수 효과, 또는 런타임 검증 가드)으로 인해 실행이 Python eager 모드로 회귀하게 되는 "그래프 브레이크(graph breaks)" 현상을 빈번하게 마주합니다.

이러한 브레이크는 계산을 여러 개의 분절된 FX 그래프로 파편화합니다. 이러한 파편화는 다음과 같은 상당한 오버헤드를 초래합니다:

1. CPU-GPU 동기화: 그래프 영역 사이에서 Python 인터프리터로 실행이 돌아오며, 이 과정에서 종종 디바이스-호스트(D2H) 메모리 전송과 동기화가 필요합니다.
2. 최적화 손실: 분절된 그래프들은 그래프 간 커널 퓨전(kernel fusion) 및 글로벌 최적화를 방해합니다.
3. 컴파일 오버헤드: 각 분절된 그래프를 별도로 추적하고 컴파일해야 하므로 콜드 스타트(cold-start) 지연 시간이 배가됩니다.
4. 서빙 제한 사항: 전체 그래프 캡처(CUDA Graphs)에 의존하는 현대적인 서빙 프레임워크(예: vLLM, SGLang)는 그래프 브레이크가 발생하는 모델을 효과적으로 활용할 수 없습니다.

195개의 Hugging Face 모델을 대상으로 한 실증 연구 결과, 13.8%(27개 모델)가 그래프 브레이크를 나타냈으며, 총 147개의 브레이크 인스턴스가 발견되었습니다. 이러한 문제는 단순한 코딩 미숙의 산물이 아니라, 구조적이며 주요 AI 조직의 프로덕션급 모델에서도 흔히 나타나는 현상입니다.

방법론: GraphMend

GraphMend는 소스 레벨 프로그램 분석 및 변환을 통해 수정 가능한 그래프 브레이크를 자동으로 제거하도록 설계된 컴파일러 기술입니다. 이는 표준 Python 소스 코드를 네이티브 바이트코드로 컴파일하는 Jaseci 프레임워크 내에서 동작합니다.

핵심 아키텍처

GraphMd는 바이트코드가 생성되기 전의 컴파일 파이프라인에 변환 단계를 삽입합니다. 이는 추상 구문 트리(AST), 제어 흐괄(CFG), 그리고 심볼 테이블(Symbol Table)을 통합한 통일된 중간 표현(UniiR)을 활용합니다. 이러한 고수준 컨텍스트를 통해 GraphMend는 코드가 바이트코드로 낮아질 때 소실되는 시맨틱(semantic) 분석을 수행할 수 있습니다.

이 기술은 두 가지 주요 패스로 구성됩니다:

1. 분석 패스(Analysis Pass): torch.compile 진입점을 식별하고 소스 레벨 패턴을 기반으로 그래프 브레이크 후보를 분류합니다.
2. AST 변환 패스(AST Transformation Pass): 브레이크를 제거하기 위한 특정 재작성 규칙을 적용합니다. 각 변환은 원래의 코드와 동일한 관측 가능한 동작을 생성함을 보장하기 위해 UniiR 상에서의 보수적인 적법성 분석(legality analysis)에 의해 보호됩니다.

세 가지 핵심 변환

GraphMend는 세 가지 범주의 그래프 브레이크를 해결합니다:

1. 조건부 데이터 의존적 제어 흐름 (Predicated Data-Dependent Control Flow):

   • 문제: 조건문(예: if x.sum() > 0)에서 조건이 런타임 텐서 값에 의존하는 경우입니다.
   • 해결책: 조건문을 torch.where 표현식으로 재작성합니다. 이를 통해 양쪽 실행 경로를 모두 FX 그래프 내에 유지하며, 텐서 값 기반의 술어(predicate)에 따라 결과를 선택합니다.
   • 제약 사항: 변환 과정에서 설정 문구(setup statements)를 호이스팅(hoisting)하며, 브래치 내부의 함수 호출이 선택되지 않은 경로에서 부수 효과나 예외를 발생시키지 않도록 보장합니다.

2. 그래프 에필로그 지연 부수 효과 (Graph-Epilogue Deferred Side Effects):

   • 문제: Python 런타임과 상호작용하는 print() 또는 로거 호출과 같은 연산입니다.
   • 해결책: 트레이싱 중에는 이러한 호출을 버퍼링하고, 그래프 실행이 완료된 후 플러시(flush)합니다.
   • 제약 사항: 호출 대상은 빌트인(built-in) 또는 지원되는 로거여야 하며, 호출의 반환 값은 사용되지 않아야 합니다. 텐서 인자는 값의 변이를 방지하기 위해 클론(clone)됩니다.

3. 조건부 트랩 로워링 (Predicated Trap Lowering):

   • 문제: 트레이싱을 깨뜨리는 검증 가드(예: if not condition: raise ValueError)입니다.
   • 해결책: Python 가드를 그래프 네이티브 단언(assertion)인 torch._assert_async로 교체합니다.
   • 메커니즘: 불필요한 텐서 비교를 피하기 위해 먼저 정적 메타데이터 체크(shape/dtype)를 수행합니다. 메타데이터가 일치하면 조건을 텐서 연산 (a == b).all()로 변환하여 비동기 단언에 전달합니다. 만약 런타임에 단언이 실패하면, 예외 핸들러가 원래의 예외 유형과 메시지를 다시 발생시킵니다.

범위 및 한계

GraphMend는 다음과 같이 소스 레벨 재작성을 통해 해결할 수 없는 브레이크는 명시적으로 제외합니다:

• tensor.item() 호출 (Python 스칼라 추출).
• 출력 형상이 런타임 값에 의존하는 동적 형상 연산자 (예: torch.nonzero).
  이들은 텐서 수준의 등가물이 없는 런타임 값을 필요로 하며, 본 도구의 범위를 벗어납니다.

주요 기여

1. 실증 연구: 195개의 Hugging Face 모델을 대규모 분석하여 그래프 브레이크의 유병률(모델의 13.8%, 브레이크 인스턴스의 73%가 소스 변환을 통해 수정 가능)을 정량화했습니다.
2. 소스 레벨 분석: AST, CFG 및 심볼 테이블을 사용하여 시맨틱 보존 변환을 위한 보수적인 적법성 조건을 확립함으로써, 바이트코드 레벨 분석에서 발생하는 시맨틱 컨텍스트의 손실을 방지하는 방법론을 제시했습니다.
3. GraphMend 기술: 바이트코드 생성 전 단계에서 그래프 브레이크를 제거하기 위해 세 가지 특정 AST 레벨 변환을 적용하는 자동화된 컴파일 패스를 구현하였으며, 이는 표준 PyTorch 2 파이프라인과 호환됩니다.
4. 포괄적 평가: NVIDIA GPU(RTX 3090, A40, H100) 환경에서 27개 모델에 대한 엄격한 벤치마크를 수행하여 출력 정확도와 성능 영향을 검증했습니다.

결과

그래프 브레이크가 존재하는 27개 모델에 대한 평가 결과는 다음과 같습니다:

• 수정률 (Fix Rate): GraphMend는 21개 모델에서 모든 그래프 브레이크를 완전히 제거했으며, 3개 모델에서는 부분적으로 수정하여 총 147개의 브서 중 107개를 제거했습니다 (73%).
• 콜드 스타트 속도 향상: 여러 개의 분절된 CUDA 그래프를 트레이싱하고 기록하는 오버헤드를 제거함으로써 최대 26배 (평균 5배)의 속도 향상을 달성했습니다.
• 스테디 스테이트(Steady-State) 속도 향상: CPU-GPU 동기화를 제거하고 커널 퓨전을 활성화함으로써 포워드 패스 레이턴시에서 최대 1.39배의 속도 향상을 달았습니다.
• 처리량 (Throughput): 모델의 추론 특성(예: 자기회귀 vs 비자기회귀)에 따라 차이가 있으나, 추론 처리량(tokens/sec 또는 samples/sec)에서 최대 **15%**의 개선을 관찰했습니다.
• 정확성: 모든 변환된 모델은 원본 모델과 비교했을 때 비트 단위로 동일한 FP32 출력과 동일한 그리디 디코딩(greedy-decoded) 토큰 시퀀스를 생성했습니다.
• 서빙 호환성: 이 기술을 통해 기존에 실패하던 모델들이 vLLM 및 SGLang와 같은 프레임워크의 필수 요구 사항인 torch.compile(fullgraph=True)를 통과할 수 있게 되었습니다.

의의

본 논문은 GraphMend를 PyTorch 2의 동적 JIT 컴파일 파이프라인을 보완하는 효과적인 도구로 포지셔닝합니다. 소스 레벨에서 그래프 브레이크를 해결함으로써, 제어 흐름, 데이터 의존성 및 부수 효과에 대한 추론의 복잡성으로 인해 발생하는 수동 개발자 리팩토링의 어려움을 해결합니다. 결과는 시맨틱을 인지하는 소스 레벨 분석이 프로덕션 환경에서 PyTorch 2 프로그램의 사용성(수동 수정 필요성 감소)과 성능(동기화 및 컴파일 오버헤드 감소)을 실질적으로 향상시킬 수 있음을 입증합니다.