vLLM-torch.compile 통합 디버깅하기

vLLM-torch.compile 통합 디버깅하기 (How to debug the vLLM-torch.compile integration)

한 줄 요약(TL;DR):

  • tlparse로 torch.compile 로그를 확보하세요. 버그 리포트·지원 요청에 이 로그를 포함하세요.
  • vLLM-torch.compile 통합은 여러 조각으로 이뤄집니다. vLLM은 각 조각을 끌 수 있는 플래그를 제공합니다:
온라인 플래그 오프라인 플래그 결과
--enforce-eager enforce_eager=True torch.compile과 CUDAGraphs 둘 다 끔
-cc.mode=0 compilation_config=CompilationConfig(mode=CompilationMode.NONE) torch.compile만 끔
-cc.mode=1 compilation_config=CompilationConfig(mode=CompilationMode.STOCK_TORCH_COMPILE) vLLM-compile이 torch.compile에 가한 수정을 끔
-cc.cudagraph_mode=NONE compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE) CUDAGraphs만 끔
-cc.backend=eager compilation_config=CompilationConfig(backend='eager') TorchInductor를 끔
-cc.ir_enable_torch_wrap=False compilation_config=CompilationConfig(ir_enable_torch_wrap=False) vLLM IR 래핑을 끔

출처: 문서

본문

vLLM-torch.compile 개요 (vLLM-torch.compile overview)

성능을 높이기 위해 vLLM은 torch.compile과 CUDAGraphs를 활용합니다. torch.compile은 PyTorch 코드에 최적화된 커널을 생성하고 CUDAGraphs는 오버헤드를 제거합니다. 가장 중요한 점은, vLLM-compile은 torch.compile이 아니라 내부 PyTorch Compile API로 구축한 커스텀 컴파일러라는 것입니다.

  • 모델이 주어지면 TorchDynamo로 배치 크기(토큰 수)에 동적인 전체 graph를 캡처합니다.
  • vLLM은 이 graph를 선택적으로 분할·특수화한 뒤 TorchInductor로 각 graph를 컴파일 산출물로 컴파일합니다. 이 단계에서 vLLM 커스텀 Inductor 패스를 사용해 graph를 더 최적화할 수 있습니다. dispatch 오버헤드를 없애기 위한 vLLM IR lowering도 포함됩니다.
  • 컴파일 산출물은 vLLM의 compile 캐시에 저장되어 나중에 로드할 수 있습니다.
  • vLLM은 CUDAGraphs를 적용해 CPU 오버헤드를 줄입니다.

네 단계 각각에서 문제가 발생할 수 있습니다. 문제가 생기면 잘못된 하위 시스템을 격리하세요. 그러면 신뢰성 목표를 유지하면서 성능 영향은 최소화하는 최소한의 것만 끄고, 버그 리포트를 열 때도 vLLM 팀에 도움이 됩니다.

설계에 대한 자세한 내용은 다음 자료를 참고하세요:

tlparse 사용 (Use tlparse)

torch.compile 로그를 보는 데 tlparse를 사용하세요. 이 로그는 컴파일 과정의 모든 단계와 torch.compile이 만드는 융합 커널을 보여줍니다.

tlparse 설치:

pip install tlparse

torch.compile 로그를 활성화하려면 환경 변수 TORCH_TRACE=<dir>를 설정할 수 있습니다. 트레이싱 동안 디렉터리 안에 랭크당 파일 하나가 생성되며, 각 파일은 컴파일 중 산출물을 담습니다. 가능하면 버그 리포트에 이 로그 파일을 포함하길 권장합니다 — 매우 유용합니다.

오프라인 추론 사용법:

TORCH_TRACE=~/trace_dir python my_script.py
tlparse ~/trace_dir/<rank_0_log_file>

서빙 사용법:

TORCH_TRACE=~/trace_dir vllm serve
# ctrl-c out of the server
tlparse ~/trace_dir/<rank_0_log_file>

로그 파일 하나가 주어지면 tlparse 명령은 HTML 파일들(예: ./tl_out/index.html)을 출력합니다. 열어서 로그를 보세요.

vLLM-torch.compile 통합 끄기 (Turn off vLLM-torch.compile integration)

--enforce-eager를 전달해 vLLM-torch.compile 통합을 끄고 전적으로 eager 모드로 실행합니다. CUDAGraphs도 함께 꺼집니다.

# Online
vllm serve --enforce-eager
# Offline
LLM(model, enforce_eager=True)

torch.compile만 끄려면 컴파일 구성에 mode = NONE을 전달합니다(-cc--compilation_config의 줄임말):

# Online
vllm serve -cc.mode=0
# Offline
from vllm.config.compilation import CompilationConfig, CompilationMode
LLM(model, compilation_config=CompilationConfig(mode=CompilationMode.NONE))

CUDAGraphs만 끄려면 cudagraph_mode = NONE을 전달합니다:

# Online
vllm serve -cc.cudagraph_mode=NONE
# Offline
from vllm.config.compilation import CompilationConfig, CUDAGraphMode
LLM(model, compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE))

vLLM IR은 functionalization, 커스텀 퓨전, lowering 등 컴파일 파이프라인을 많이 사용합니다. 그것을 끄고 vLLM IR의 eager 모드 디스패칭 동작을 캡처하려면 ir_enable_torch_wrap=False로 실행하세요. IR torch wrap은 mode=VLLM_COMPILEbackend="inductor"(기본)를 쓸 때만 기본 활성화됩니다.

# Online
vllm serve -cc.ir_enable_torch_wrap=False
# Offline
from vllm.config.compilation import CompilationConfig
LLM(model, compilation_config=CompilationConfig(ir_enable_torch_wrap=False))

TorchDynamo 디버깅 (Debugging TorchDynamo)

vLLM은 모델 코드가 TorchDynamo(torch.compile의 프론트엔드)를 통해 전체 graph로 캡처 가능해야 합니다. TorchDynamo는 모든 Python을 지원하지 않습니다. 지원하지 않는 기능을 만나면 (fullgraph 모드에서) 오류를 냅니다 (이를 graph break라고도 합니다).

graph break를 만나면 pytorch/pytorch에 이슈를 열어 PyTorch 개발자가 우선순위를 둘 수 있게 하세요. 그다음 graph break를 피하도록 코드를 최대한 다시 작성하세요. 자세한 내용은 이 Dynamo 가이드를 참고하세요.

동적 형태(Dynamic Shape) 전체 graph 캡처 디버깅 (Debugging Dynamic Shape full graph capture)

vLLM은 모델의 forward pass가 배치 크기(즉 토큰 수)에 동적인 전체 graph로 캡처 가능해야 합니다. (기본적으로) 이 하나의 graph를 하나의 산출물로 컴파일하고 모든 배치 크기에 그 산출물을 사용합니다.

코드가 Dynamic Shapes로 캡처될 수 없으면 조용한 부정확성, 큰 오류, 또는 CUDA 불법 메모리 접근이 나타날 수 있습니다. 예를 들어 다음은 단일 graph로 캡처할 수 없습니다:

if data.size[0] % 128 == 0:
    foo(...)
else:
    bar(...)

이 문제는 진단하기 쉽습니다. tlparse를 쓰고 compilation_metrics를 클릭하면 배치 크기에 대한 기호 제약(symbolic constraints)이 표시됩니다. 배치 크기를 제한하는 제약이 있으면 문제가 있는 것입니다.

이를 피하려면:

  • 토큰 수에 분기하지 않거나
  • 분기 로직을 커스텀 연산자로 감싸세요. TorchDynamo는 커스텀 연산자를 트레이싱하지 않습니다.

제약 위반 및 동적 형태 가드 문제 디버깅 (Debugging constraint violations and dynamic shapes guards issues)

동적 형태 가드(dynamic-shape guards)는 Dynamo 가드의 특정 범주입니다. torch.compile이 동적 차원(예: seq_len)에 부착해 컴파일 산출물이 유효하게 유지되도록 하는 제약입니다. 이 가드는 보통 프레임워크 코드, 커스텀 패스, 또는 사용자 코드가 동적 형태 값에 따라 분기할 때 나타납니다.

예:

if x > 10:
    # path A
else:
    # path B

이것은 트레이싱된 경로에 따라 x > 10 또는 x <= 10 가드를 만듭니다.

vLLM의 가정: vLLM은 torch.compile이 추가한 모든 가드를 버려도 안전하며 컴파일된 graph를 특정 입력 형태로 제약하지 않을 것이라고 가정합니다. 이 가정이 위반되면 사용자가 디버깅해야 하는 문제가 발생합니다. 이 가정이 위반됐다는 부작용으로는 런타임 오류나 ConstraintViolationErrors가 있습니다.

동적 형태가 단일 값으로 제약되면 ConstraintViolationErrors가 발생합니다. 제약 위반 오류를 만나거나 동적 형태 가드가 잘못 추가되고 있다고 의심되면 더 엄격한 동적 형태 모드를 사용해 문제를 격리할 수 있습니다:

# Online - using unbacked mode
vllm serve meta-llama/Llama-3.2-1B -cc.dynamic_shapes_config.type=unbacked

# Online - using backed_size_oblivious mode
vllm serve meta-llama/Llama-3.2-1B -cc.dynamic_shapes_config.type=backed_size_oblivious
# Offline - using unbacked mode
from vllm.config.compilation import CompilationConfig, DynamicShapesConfig, DynamicShapesType
LLM(model, compilation_config=CompilationConfig(
    dynamic_shapes_config=DynamicShapesConfig(type=DynamicShapesType.UNBACKED)
))

# Offline - using backed_size_oblivious mode
from vllm.config.compilation import CompilationConfig, DynamicShapesConfig, DynamicShapesType
LLM(model, compilation_config=CompilationConfig(
    dynamic_shapes_config=DynamicShapesConfig(type=DynamicShapesType.BACKED_SIZE_OBLIVIOUS)
))

이 모드들은 더 엄격하고 동적 형태 가드의 필요를 줄이거나 없애 문제를 격리하는 데 도움이 됩니다:

  • unbacked: 가드를 허용하지 않는 unbacked symint를 사용해 가드가 어디에 잘못 추가되는지 더 쉽게 식별합니다.
  • backed_size_oblivious: 가드에 대해 더 엄격한 모드를 사용합니다.

동적 형태 모드에 대한 자세한 내용은 동적 형태와 vLLM 가드 드롭을 참고하세요.

가드 출력 (Printing guards)

컴파일 중 추가되는 모든 가드를 보려면 TORCH_LOGS=+dynamic을 사용하세요:

TORCH_LOGS=+dynamic vllm serve meta-llama/Llama-3.2-1B

로그에서 [guard added]를 찾아 어디에 가드가 추가되는지 확인하세요. 어떤 연산이 가드를 잘못 추가하는지 식별하는 데 도움이 됩니다.

TorchInductor 디버깅 (Debugging TorchInductor)

TorchInductor는 캡처된 graph를 가져와 1개 이상의 triton 커널을 호출할 수 있는 Python 코드로 컴파일합니다. 드물지만 불행한 경우 잘못된 triton 커널을 만들 수 있습니다. 이는 조용한 부정확성, CUDA 불법 메모리 접근, 또는 큰 오류로 나타날 수 있습니다.

Inductor 런타임 어서션 (Inductor runtime assertions)

기본적으로 (torch < 2.12에서) vLLM은 대형 모델의 forward pass당 약 2ms 오버헤드를 피하기 위해 Inductor의 런타임 어서션(assert_size_stride, assert_alignment)을 비활성화합니다. VLLM_LOGGING_LEVEL=DEBUG를 설정하면 자동으로 다시 활성화돼 디버깅 세션에서 완전한 형태/스트라이드 검증을 얻습니다:

VLLM_LOGGING_LEVEL=DEBUG vllm serve <model>

--compilation-config로 명시적으로 덮어쓸 수도 있습니다:

vllm serve <model> -cc.inductor_compile_config='{"size_asserts": true, "alignment_asserts": true, "scalar_asserts": true}'

torch >= 2.12에서는 PyTorch가 효율적인 assert-once 전략을 사용해 이 플래그들이 더 이상 vLLM에 의해 억제되지 않습니다.

TorchInductor가 문제인지 디버깅하려면 컴파일 구성에 backend='eager'를 전달해 비활성화할 수 있습니다:

# online
vllm serve -cc.backend=eager
# offline
LLM(compilation_config=CompilationConfig(backend='eager'))

Inductor가 문제라면 PyTorch에 버그를 제기하세요. 모험심이 있다면 tlparse로 찾을 수 있는 Inductor 출력 코드의 triton 커널을 디버깅할 수 있습니다.

또한 TORCH_LOGS=output_code <command>를 사용해 Inductor 출력 코드를 출력할 수 있습니다.

편집 가능한 TorchInductor 코드 (Editable TorchInductor code)

VLLM_COMPILE_CACHE_SAVE_FORMAT=unpacked를 설정하거나 -cc.compile_cache_save_format=unpacked를 전달해 실행되는 TorchInductor 코드를 편집할 수 있습니다. 기본값은 binary로 편집할 수 없습니다.

이는 유용한 기법입니다. 출력 코드에 중단점(예: torch.distributed.breakpoint())과 print 문을 넣을 수 있습니다.

vLLM-compile 캐시 디버깅 (Debugging vLLM-compile cache)

vLLM은 torch.compile 산출물을 위한 자체 캐시를 구축했습니다. 아이디어는 산출물을 한 번 컴파일한 뒤 재사용하는 것입니다. 이는 torch.compile의 컴파일러 캐시 위의 계층입니다.

torch.compile의 컴파일러 캐시는 매우 안정적인 반면, vLLM의 컴파일러 캐시는 불행히도 항상 올바르지는 않습니다. VLLM_DISABLE_COMPILE_CACHE=1을 설정해 비활성화할 수 있습니다.

또한 이 캐시를 수동으로 제거할 수 있습니다.

  • rm -rf ~/.cache/vllm로 vLLM의 compile 캐시를 제거합니다 (로그에서 위치가 바뀌었는지 확인).
  • rm -rf /tmp/torchinductor_$(whoami)로 torch.compile의 내장 캐시를 제거합니다.

vLLM의 캐시는 캐시 키→컴파일 산출물 매핑입니다. vLLM은 여러 요인(예: 구성 플래그와 모델 이름)을 결합해 캐시 키를 계산합니다. vLLM의 compile 캐시가 틀리다면 보통 어떤 요인이 빠진 것입니다. vLLM이 캐시 키의 일부를 계산하는 방법은 이 예제를 참고하세요.

vLLM의 컴파일 캐시는 컴파일되는 코드가 직렬화 가능해야 합니다. 그렇지 않으면 저장 시 오류가 납니다. 보통 해결책은:

  • 직렬화 불가능한 조각을 다시 작성 (지금 무엇이 직렬화 가능하고 무엇이 아닌지 판별하기 어려워 어려울 수 있음)
  • 버그 리포트 제기
  • VLLM_DISABLE_COMPILE_CACHE=1을 설정해 오류 무시 (이 경우 워밍 서버 시작이 훨씬 느려집니다)

CUDAGraphs 디버깅 (Debugging CUDAGraphs)

CUDAGraphs는 다음을 가능하게 하는 기능입니다:

  • 1개 이상의 CUDA 커널을 실행하는 callable을 CUDAGraph로 캡처
  • CUDAGraph 재생(replay)

캡처된 CUDAGraph는 캡처 과정에서 사용된 모든 메모리를 포함합니다. CUDAGraph의 재생은 정확히 같은 메모리 영역을 읽고 씁니다.

이로 인해 몇 가지 제약이 있습니다:

  • 새 데이터에 CUDAGraphs를 사용하려면 CUDAGraph가 읽는 버퍼에 데이터를 복사해야 합니다.
  • CUDAGraphs는 CUDA 커널만 캡처하며 CPU에서 수행되는 작업은 캡처하지 않습니다.

vLLM은 원시 CUDAGraphs API를 사용하며, 잘못 사용하면 안전하지 않습니다.

CUDAGraphs만 끄려면 cudagraph_mode = NONE을 전달합니다:

# Online
vllm serve -cc.cudagraph_mode=NONE
# Offline
from vllm.config.compilation import CompilationConfig, CUDAGraphMode
LLM(model, compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE))

더 알아보기 (Learn more)