PyTorch - A Reference Language

이 글은 PyTorch: a reference language를 참고하여 작성했습니다.

7월 25일 PyTorch devlog에 Edward Z. Yang(ezyang)이 글을 하나 올렸습니다. 제목은 "PyTorch: a reference language"입니다.

주장은 이렇습니다. PyTorch는 지금 두 역할을 겸하고 있습니다. 하나는 참조 언어(reference language), 다른 하나는 구현 언어(implementation language)입니다. 규모가 커질수록 이 둘을 갈라 놓고, 참조 구현을 프로덕션 구현의 정확성을 검증하는 별도 산출물로 보는 편이 자연스러워집니다.

Yang이 요약한 문장이 정확합니다. 연구할 구현 하나, 스케일할 구현 하나, 그리고 둘을 묶는 검증자 하나.

참조와 구현

Yang은 통념의 모순을 먼저 짚습니다.

참조 구현은 보통 프로덕션에 나가지 않습니다. 그런데 사람들은 PyTorch로 학습을 돌립니다. 성능이 중요한 연산은 커널 DSL이 맡고 있습니다. 그리고 AI 코딩이 있으면 원칙적으로 스택 전체를 처음부터 다시 쓸 수도 있습니다.

이 긴장을 푸는 방식이 역할 분리입니다. 규모가 작거나 컴파일러가 잘 작동할 때는 참조 구현을 그대로 배포해도 됩니다. 그렇지 않을 때는 참조 구현이 배포 대상이 아니라 검증 대상이 됩니다.

컴파일러가 알아서 해 준다는 약속

이 글이 눈길을 끄는 이유는 PyTorch 코어 개발자가 오랜 전제를 공개적으로 조정하고 있기 때문입니다. 계보를 짚으면 선명합니다.

초기 PyTorch의 정체성은 eager mode였습니다. 파이썬 코드가 쓰인 순서대로 실행되고, 디버깅이 되고, 그래프를 미리 만들 필요가 없었습니다. 대신 최적화 여지를 포기했습니다.

TorchScript는 그 그래프를 되찾으려는 시도였습니다. 파이썬의 부분집합을 정적으로 뽑아 컴파일하는 방식이었는데, 파이썬의 동적 특성 때문에 커버리지가 늘 문제였습니다.

torch.compile은 접근을 바꿨습니다. 바이트코드를 가로채 그래프를 뽑고, 안 되는 부분은 파이썬으로 되돌리는(graph break) 구조입니다. 커버리지 문제를 실용적으로 우회했고, 여기서 "고수준 PyTorch를 쓰면 컴파일러가 알아서 빠르게 만들어 준다"는 약속이 가장 강했습니다.

이번 글은 그 약속의 범위를 좁힙니다. 실제로 벌어진 일은 컴파일러가 자동 최적화를 완성한 것이 아니라, 성능이 중요한 연산을 사람이 커널 DSL로 직접 쓰는 관행이 자리 잡은 것입니다. 행렬 곱, 어텐션 같은 것들입니다.

커널 저자들이 이미 하고 있는 것

Yang이 근거로 드는 관찰이 중요합니다. 커널 DSL로 최적화 커널을 쓰는 사람들 대부분이 PyTorch 참조 구현을 나란히 유지합니다. 그리고 수치 비교로 정확성을 검증합니다.

즉 참조 구현과 최적화 구현의 이원 구조는 제안이 아니라 이미 있는 관행입니다. 이 글은 그 관행에 이름을 붙이고 범위를 넓히자고 말합니다.

여기서 참조 구현의 가치가 실행 속도가 아니라는 점이 드러납니다. 아무도 참조 커널로 프로덕션을 돌리지 않습니다. 가치는 "이 연산이 정확히 무엇을 계산하는가"를 읽을 수 있게 적어 둔 것에 있습니다. 실행 가능한 명세입니다.

autograd가 걸림돌이 되는 지점

프로덕션 규모에서는 PyTorch의 자동 미분이 문제가 된다는 것이 다음 논점입니다.

암묵적 backward 그래프가 짐이 됩니다. 계산의 대부분이 눈에 보이지 않는 곳에서 일어나고, 일반적인 디버깅이 잘 통하지 않고, 융합(fusion) 기회를 잡기도 어렵습니다. Yang은 소스-투-소스 자동 미분을 탐색했던 Google Tangent를 언급합니다. 지금은 없어진 라이브러리입니다.

forward만 명시적이고 backward가 프레임워크 안에 숨어 있으면, 학습 스텝 전체를 최적화 대상으로 다루기 어렵습니다. 커널 하나는 고칠 수 있어도 학습 루프 전체는 그렇지 않습니다.

LLM이 끼어드는 자리

Yang이 제안하는 패턴은 이렇습니다. PyTorch autograd 코드는 참조 명세로 남깁니다. 그리고 명시적인 forward-backward 구현을 LLM으로 생성해 따로 최적화합니다. 컴파일러의 취약한 패턴 매칭에 기대지 않는 방식입니다.

대가는 두 버전이 갈라질 수 있다는 점입니다. 그래서 검증 장치가 필요합니다. 비트 단위 동치 테스트를 쓰거나, 그래프 구조를 비교하는 방식(translation validation)을 씁니다. 특정 융합에 대한 참조 구현을 따로 두어 차이를 흡수하는 방법도 포함됩니다.

이 대목이 이 글에서 가장 실무적으로 중요합니다. 사람이 최적화 커널을 쓰는 대신 LLM이 쓰게 되면, 사람이 하던 일이 "코드 작성"에서 "채점 기준 관리"로 옮겨갑니다. 그러면 참조 구현이 인간의 산출물 중 가장 값비싼 것이 됩니다. 명세가 틀리면 생성된 코드 전부가 틀립니다.

실무 함의

"최적화 코드를 쓸 때 참조 구현을 함께 두세요"

이미 커널을 쓰는 사람이라면 하고 있을 겁니다. 그렇지 않은 경우에도 적용됩니다. 성능을 위해 뭔가를 손으로 풀어 쓸 때, 느리지만 읽으면 명확한 버전을 옆에 남기고 수치 비교 테스트를 붙여 두는 것입니다. 그 버전이 나중에 리팩터링의 기준이 되고, LLM에게 최적화를 시킬 때 채점표가 됩니다.

Yang 본인은 이 방식이 모두에게 맞는다고 주장하지 않습니다. PyTorch가 여전히 꽤 좋은 실행 가능한 명세라는 점을 인정하며, 이 관점을 Horace He의 질문에 대한 답으로 위치시킵니다. eager mode의 제어권과 그래프 수준 추상화의 편의를 어떻게 함께 가져갈 것인가라는 질문이었습니다.

정리

컴파일러가 고수준 코드를 자동으로 최적화한다는 약속은 부분적으로만 실현됐고, 실제로는 커널 DSL이 성능을 맡고 PyTorch가 명세를 맡는 분업이 자리 잡았습니다. Yang은 그 분업을 인정하고 이름을 붙입니다.

여기에 LLM이 들어오면 분업의 무게중심이 이동합니다. 최적화 코드는 생성물이 되고, 참조 구현은 그 생성물을 채점하는 기준이 됩니다. 프레임워크의 가치가 실행 성능에서 명세의 명확성으로 옮겨간다는 뜻입니다.

같은 주장을 다른 자리에서 하는 작업이 최근 몇 건 있었습니다. 07-22에 나온 Molt 논문은 에이전트 RL 프레임워크의 설계 기준으로 "코딩 에이전트가 코드베이스 전체를 읽고 추론할 수 있는가"를 명시적으로 넣었습니다. LLM 가독성이 프레임워크 설계 기준이 되는 흐름이 서로 다른 곳에서 동시에 나타나고 있습니다.

이 글은 Edward Z. Yang의 PyTorch devlog의 관점에서 작성되었습니다.