-
Notifications
You must be signed in to change notification settings - Fork 2
2회차 받아쓰기
시작하겠습니다. 이 페이지는 상단 링크에 있는 GitHub 페이지의 내용을 가져와 정리한 것입니다. 전체적으로 세 개의 섹션으로 나뉘어 있는데요. 첫 번째는 인트로, 두 번째는 PyTorch의 기초적인 소개, 그리고 마지막으로 주제별로 심화된 내용을 다루는 구성입니다. 오늘은 그중 기초 섹션의 첫 번째 시간으로, Eager 모드에 대해 자세히 이야기해 보려고 합니다.
Eager 모드의 아키텍처에 대해 개괄적으로 설명한 뒤, 실제로 torch.matmul 연산자가 수행될 때 어떤 과정이 이루어지는지 살펴보겠습니다. 시간이 허락된다면, 그 뒤에 있는 세부적인 내용도 다뤄볼 예정입니다. 다만, 준비한 자료의 양이 많아 모든 내용을 다 다루기는 어려울 것 같고, 적절한 시점에서 마무리해야 할 것 같습니다. 추가로 공부하고 싶으신 분들은 이후의 자료를 참고하여 스스로 학습을 이어가시면 좋을 것 같습니다.
아키텍처 그림을 제가 빠뜨렸네요. 일단 PyTorch에서 Python 코드를 작성하고, import torch를 통해 라이브러리를 불러온 다음, 텐서를 만들어 연산을 수행하면 결과값을 얻을 수 있죠. 그 결과를 print로 확인할 수도 있고요. 그런데 여기서 제가 특히 궁금했던 부분은, 실제로 matmul operator 같은 연산을 수행할 때 내부적으로 어떤 일이 벌어지는지였습니다.
여러분들도 이 부분을 공부해 보셨을 것 같은데, Eager 모드에서 어떤 일이 일어나는지 살펴보면 생각보다 굉장히 많은 과정이 진행됩니다. 이를 확인할 수 있는 방법 중 하나가 콜스택(call stack)을 분석하는 것입니다. 특히 matmul같은 경우에는 Eager 모드에서 어떤 과정을 거치는지 콜스택을 통해 거의 대부분 확인할 수 있습니다. 이를 통해 Pytorch 내부에서 어떤 일이 일어나는지 명확히 이해할 수 있는 좋은 사례라고 생각합니다.
이 콜스택을 얻은 방법은 다음과 같습니다. 먼저 PyTorch가 GPU를 사용하도록 설정한 뒤, CUDA의 GEMM(general matrix multiplication) 함수로 진입할 것으로 예상되는 지점에 브레이크포인트를 설정했습니다. 그런 다음, Eager 모드에서 matmul 연산을 수행하여 브레이크포인트가 걸리도록 하고, 그 상태에서 콜스택의 스냅샷을 찍은 것입니다.
여기 있는 내용을 하나하나 모두 읽으실 필요는 없습니다. 다만, 제가 보여드리고자 했던 것은, PyTorch가 연산을 수행하며 여러 중간 단계를 거쳐 최종적으로 CUDA의 GEMM 함수에 도달한다는 점입니다. 이 과정은 여기 보시는 것처럼 굉장히 복잡해 보이고, 실제로도 많은 일이 진행됩니다. 이를 통해 PyTorch 내부에서 연산이 수행되는 과정을 이해하는 데 도움이 될 것이라 생각합니다.
이게 왜 이렇게 복잡할까요? 첫째, Torch의 mm operator 자체가 복잡합니다. 단순히 2D matrix 곱셈만을 수행하는 것이 아니라 다양한 상황과 경우를 처리해야 합니다. 이로 인해 내부적으로 수행하는 작업이 많고, 최종적으로는 2D matrix 곱하기 2D matrix인 gemm(General Matrix Multiplication)으로 진입하는 경우도 있지만, 그 과정이 간단하지 않습니다. 둘째, 여러 language boundary를 넘어야 합니다. PyTorch는 단일 언어로 구현된 것이 아니라 Python, C++, CUDA 등 여러 언어가 혼합된 형태입니다. 연산이 언어 경계를 넘어갈 때마다 복잡성이 추가되며, 이러한 부분들이 call stack에도 어느 정도 반영됩니다. 셋째, matmul operator의 상황별로 다른 역할을 하기 때문입니다. matmul operator는 실행 환경이나 설정에 따라 매우 다른 역할을 합니다. CPU, GPU, NPU 등 다양한 디바이스에서 연산이 수행될 수 있습니다. 연산이 forward 계산인지, 혹은 backward propagation까지 포함하는 training 과정인지에 따라 dynamic computation graph가 생성될 수도, 생성되지 않을 수도 있습니다. PyTorch 1.0의 경우, graph mode를 위해 trace 기능이 사용될 수도 있고, 아닐 수도 있습니다. 이러한 다양한 실행 시나리오로 인해 같은 matul operator라도 내부 동작이 달라질 수 있습니다.
오늘 다룰 내용은 지난주에 정리했던 PyTorch의 주요 특성들과 연결됩니다. 먼저, Numpy-like 경험에 주목해보면 PyTorch는 interactive mode를 통해 Tensor와 다양한 operator를 활용해 즉각적으로 결과를 확인하며 프로그래밍 및 데이터 분석이 가능하다는 점이 특징입니다. 이러한 사용자 경험을 제공하면서도 내부적으로는 복잡한 연산 처리를 지원해야 합니다. 다음으로, Heterogeneous Computing과 관련해 PyTorch는 CPU, GPU, NPU 등 다양한 디바이스를 지원합니다. 동일한 operator라도 실행되는 디바이스에 따라 동작 방식이 달라지기 때문에 call stack에 이러한 복잡성이 반영됩니다. 여기에 더해, PyTorch는 Python, C++, CUDA로 구성된 Three Language Layers를 가지고 있어 언어 경계를 넘어가는 과정에서도 추가적인 복잡성이 발생합니다.
또한, ML compiler와 compute library와의 통합 역시 PyTorch의 중요한 특성 중 하나입니다. 코드 내부로 깊이 들어가지는 않겠지만 이러한 통합이 call stack에 어떤 영향을 주는지 확인해 볼 수 있습니다. 한편, PyTorch의 코드베이스에서는 Tool-based Code Generation이 중요한 역할을 합니다. 반복적이고 비슷한 기능을 하는 코드를 수작업 대신 tool을 통해 생성하도록 설계되어 있으며, 특히 autograd 구현에서 이러한 방식이 활용됩니다. 이러한 자동화된 코드 생성 방식은 코드베이스를 더욱 복잡하게 만듭니다. 마지막으로, 같은 operator라도 실행 환경이나 모드에 따라 수행하는 작업이 달라집니다. 예를 들어, forward 계산과 backward propagation 여부, dynamic computation graph의 생성 유무 등에 따라 다른 동작을 하게 됩니다. 이제 이러한 특성들이 call stack에 어떻게 반영되어 있는지를 분석하며 구체적으로 살펴보겠습니다.
오늘 살펴볼 call stack은 아래에서 위로 코드가 진입한다고 생각하시면 됩니다. 이를 시각적으로 표현하기 위해 화살표를 밑에서 위로 표시했습니다. 가장 아래에 위치한 부분은 CPython이 수행하는 call stack입니다. CPython은 C나 C++로 작성된 프로그램이기 때문에 자체적으로 call stack이 존재하며, Python 코드가 실행되는 기반이 됩니다. 다만, 이 부분은 오늘 논의에서 제외하고 넘어가겠습니다. 위로 올라가면 PyTorch 내부 코드에 해당하는 부분을 확인할 수 있습니다. 이 영역에서 PyTorch의 다양한 연산과 operator들이 어떻게 동작하는지를 집중적으로 분석할 것입니다. 반면, call stack의 맨 위에 있는 부분은 오늘 다룰 내용과 크게 관련이 없으므로 신경 쓰지 않아도 됩니다. 이제 PyTorch 내부 영역을 중심으로 call stack을 단계별로 분석해 보겠습니다.
네 이제 앞으로 이제 조금 더 자세히 코드를 보기 시작할 텐데 일단 이걸 잠깐 언급을 하고 넘어가고 싶었어요. 그러니까 전체적으로 eager mode가 어떻게 생겼냐면 실제 op을 수행을 하면 그 op이 뭔지 파악을 하고 Python에서 시작해서 C++로 넘어오는 약간 front-end에 해당하는 그런 부분이 있고요. 그다음에 그게 실제 어떤 op인지라고 하는 게 판별이 되면 근데 그게 어떤 일을 해야 되는가라고 하는 거를 판별을 해서 그 일을 하는 backend를 이제 끄집어내는 dispatcher라고 하는 과정이 있고 그 dispatcher가 특정 device backend에 붙어 있는 그런 op의 구현을 불러내는데 거기서 dispatch가 한 번 되고 끝나는 게 아니라 re-dispatch라고 하는 그런 과정이 있을 수가 있어요. 그리고 이 re-dispatch가 생각보다 여러 번 반복이 되는 경우들도 있거든요. 뒤에서 보시겠지만 이 call stack의 depth가 굉장히 길어진 가장 중요한 이유가 dispatch와 그러니까 re-dispatch가 반복이 되기 때문에 생기게 된 거고요. device backend는 실제 device에 이제 올라가서 뭔가 일을 하기 위해서 그 device의 runtime을 타고 들어가게 돼 있습니다.
이제 PyTorch eager mode가 동작하는 원리를 high-level에서 이해하고 나서, 다음 슬라이드들을 보시면 더 좋을 것 같습니다. 두 번째로 주목할 만한 특성은 자동 생성된 코드가 많다는 점입니다. 앞서 말씀드렸듯이, PyTorch 내부 코드 중 일부는 자동 생성된 코드로 구성되어 있습니다. 특히 제가 색깔을 입혀놓은 부분들이 그에 해당합니다. 이 자동 생성된 코드는 주로 연산(op)별로 반복적으로 수행되는 작업을 처리하는 데 사용됩니다. PyTorch에서는 이런 반복되는 작업들을 사람이 직접 작성하는 대신 code generation(codegen)을 통해 자동으로 생성하고 있습니다. 이러한 코드가 어떤 기능을 수행하는지에 대해서는 뒷부분에서 다시 한번 자세히 다룰 예정입니다. 참고로 말씀드리면, PyTorch의 code generation은 build time에 이루어집니다. PyTorch에서는 사람이 작성한 코드뿐만 아니라 build time에 자동으로 생성되는 코드들도 함께 포함됩니다. 이 모든 코드가 build 과정을 거쳐 최종적으로 binary나 Python 패키지로 만들어지게 됩니다.
여기서 제가 정리한 내용은 자동 생성되는 파일들이 무엇인지, 어떤 도구를 통해 생성되는지, 그리고 그 생성 과정에서 사용되는 입력 파일이 무엇인지를 설명한 것입니다. 먼저 native_functions.yaml은 매우 중요한 역할을 합니다. 이 파일은 Python에서 기본적으로 제공하는 연산(op) 목록과 각 연산의 특성 및 동작 방식을 기술한 스펙이라고 생각하면 됩니다. PyTorch 내부적으로 이 파일이 가장 중요한 backbone 역할을 한다고 볼 수 있습니다. 다음으로 derivatives.yaml은 backward autograd를 위한 정보를 기술한 파일입니다. 이를 통해 자동으로 미분이 가능하도록 하는 데 활용됩니다. 다시 언급된 native_functions.yaml은 동일한 파일이므로 추가 설명은 생략하겠습니다. 이러한 정보를 바탕으로 gen.py와 같은 도구들이 다양한 파일들을 생성하게 됩니다. 시간 관계상 각각의 생성된 파일을 자세히 다루기는 어렵지만, ATen에 해당하는 부분은 PyTorch의 front-end 역할을 한다고 볼 수 있습니다. 특히 register backend 부분은 중요한 초기화 과정 중 하나입니다. PyTorch는 build time에 코드가 생성되고, 이후 실제 실행 시 초기화 과정에서 다양한 backend들을 등록(registration) 합니다. 이 과정에서 자동 생성된 코드들이 backend에 따라 여러 파일로 나뉘어 등록되며, 그 코드들이 여기 표시된 부분에 포함됩니다. 이 부분에 대한 더 깊은 이해를 원하시면, 실제 PyTorch code base를 살펴보는 것이 가장 효과적일 것입니다. 오늘 다룬 내용을 기반으로 직접 코드를 확인해보면 PyTorch의 구조를 더 명확하게 파악하실 수 있을 것입니다.
Dispatch logic은 PyTorch의 중요한 특성 중 하나이며 코드에서도 명확히 드러납니다. 제가 노란색으로 표시한 부분은 코드의 세부 내용을 보시라는 의미가 아니라, 어떤 부분인지 시각적으로 강조하기 위해 색을 입힌 것입니다. Dispatch logic은 주로 템플릿을 기반으로 반복되는 구간을 구현하는 코드입니다. 이는 앞서 설명한 dispatcher가 연산(op)을 선택하고 실행하며, 필요에 따라 re-dispatch를 수행하는 과정을 담당합니다. 예를 들어, 특정 연산인 mm operator만 살펴봐도 이 과정이 세 번 반복되는 것을 확인할 수 있습니다. 이 과정에서는 연산이 실행될 때 어떤 조건에서 어떤 작업을 수행해야 하는지를 결정하는 절차가 필요합니다. 이를 위해 dispatch key와 dispatch key set이라는 개념이 PyTorch 내부에 구현되어 있습니다. 이 key set을 통해 dispatcher가 어떤 구현체를 호출해야 하는지 판단하고, 등록된 구현체를 실행하는 과정을 수행합니다. 이러한 일련의 흐름을 dispatch logic이라고 표현했습니다.
PyTorch의 dispatcher는 Meta의 개발자 Edward Yang이 설계한 것으로 보입니다. 그의 설계 철학과 구현 과정에 대한 상세한 설명은 제가 공유한 링크에 잘 나와 있습니다. 다만 현재의 구현은 그 당시와 조금 달라진 부분도 있지만, 핵심 개념은 여전히 유지되고 있습니다. Dispatcher는 C++의 vtable과 유사한 개념을 가지고 있습니다. vtable은 객체지향 프로그래밍에서 virtual function의 동작을 지원하기 위해 사용되는 구조입니다. 예를 들어, 부모 클래스에서 정의된 virtual function이 자식 클래스에서 오버라이드되면, 자식 클래스의 인스턴스는 vtable을 통해 오버라이드된 함수의 포인터를 참조하게 됩니다. 이처럼 vtable은 동일한 위치에 서로 다른 기능이 실행되도록 도와줍니다.
하지만 PyTorch에서는 vtable을 그대로 사용할 수 없었기 때문에, 더 복잡하고 다양한 context를 지원하는 별도의 데이터 구조와 알고리즘을 설계했습니다. 이 구조의 핵심은 dispatch table이며, 특정 key에 따라 해당하는 함수 포인터를 호출하는 방식으로 동작합니다. PyTorch에서는 단일 key가 아니라 key set 개념을 도입했으며, 여러 key가 동시에 의미를 가질 수 있도록 설계되었습니다. key set에서 우선순위(priority)가 가장 높은 key를 선택해 dispatcher가 적절한 구현체를 호출하게 됩니다. 청중분께서 언급하신 autograd나 CUDA처럼 서로 독립적인(orthogonal) 정보들이 동시에 적용되는 경우, key의 동작이 어떻게 결정되는지에 대한 질문은 타당합니다. 이 부분에 대한 답은 저희 팀도 고민했었고 어느 정도 이해한 상태입니다. 이 주제는 뒷부분에서 더 자세히 다룰 예정이므로 그때 다시 살펴보도록 하겠습니다.
실제 dispatcher logic과 반복되는 부분, 그리고 코드 생성된 부분들을 제외하고 보면, 개발자가 신경 써서 작성해야 하는 코드 부분은 생각보다 적다는 점을 강조하고 싶습니다. 이 부분들은 실제 mm operator의 동작과 밀접하게 연결되어 있는 핵심 부분입니다. 한 부분은 매우 일반적인 코드이고, 그 위의 부분은 CUDA에서 mm operator가 실행되기 위해 필요한 작업들을 기술한 부분입니다. 이처럼, call stack이 55개의 함수 깊이를 가지며 복잡한 작업을 수행한다고 해서 엄청나게 복잡한 구조라고 생각할 수 있지만, 실제로는 많은 부분이 코드 생성과 관련된 것이며, 반복적으로 수행되는 dispatcher logic이 많습니다. 이러한 요소들을 제외하고 나면, 실제로 mm operator와 밀접하게 연관된 함수들은 관리할 수 있을 정도로 간소화된다고 할 수 있습니다.
지금까지는 call stack이 수행하는 작업의 특성보다는 각기 다른 코드들로 구분되는 특성들을 살펴봤습니다. 이제 실제로 call stack이 어떤 일을 하는지에 대해 구체적으로 살펴보겠습니다. 맨 밑에 있는 부분은 CPython이라고 하는 Python 처리 프로그램이 실행되는 부분입니다. CPython은 일반적인 C나 C++ 프로그램처럼 libc가 main을 호출하고, 그 이후 Python 처리를 위한 여러 작업을 수행하면서 C++ 코드로 넘어가는 단계가 있습니다. 이때 C++ 코드로 넘어가는 과정은 matmul operator의 C++ 내부 구현체로 가기 위한 준비 작업이 이루어집니다. 요즘은 pybind11을 사용하여 C++와 Python을 연결하는 일이 많이 쉬워졌지만, PyTorch에서는 이를 사용하지 않고 커스텀 방식으로 바인딩을 구현한다고 합니다. 그 이유는 성능 문제와 코드 생성(codegen)과의 연관, 그리고 pybind11을 사용하기 전에 이미 시작된 구현 방식 때문일 수 있습니다.
C++ 코드로 넘어가면 먼저 front-end 역할을 하는 일반적인 작업이 이루어지고, 이후 mm operator kernel을 찾아서 디스패치(dispatch)하는 과정이 시작됩니다. 제가 처음에 설명했듯이 matmul 오퍼레이터는 다양한 케이스를 처리하기 때문에, 이 함수는 실제로 여러 작업을 수행합니다. 그 작업을 진행하면서 matmul operator는 2D 행렬을 계산하는 작업을 다시 한 번 디스패치하게 되며, 이때 autograd kernel을 호출하게 됩니다. 이 부분은 제가 코드를 직접 확인하지는 않았지만, 이 커널이 autograd key에 해당하는 커널을 디스패치한다고 추측됩니다. 이 과정은 청중이 질문하신 부분과 관련이 있는데, 그에 대한 설명은 조금 더 뒤에서 다루겠습니다. 이후 autograd kernel이 실제 계산을 진행하면서, 결국에는 cublas로 구현된 gemm API 중 하나를 호출하는 형태로 CUDA 코드로 진행되어 최종적으로 gemm 코드가 실행됩니다.
이 call stack이 실제로 수행하는 작업은 위에서 설명한 대로 PyTorch가 Python에서 C++로 넘어가는 과정입니다. PyTorch는 이 과정에서 pybind11을 사용하지 않고, 성능과 확장성의 요구를 충족시키기 위해 custom binding 기술을 사용한다고 합니다. 이 custom binding 구현체는 자동화된 코드 생성 및 dispatcher 메커니즘과 잘 맞물려 있습니다. 이와 관련된 더 구체적인 설명은 위 링크를 참조하면 좋을 것 같습니다. 또한, 이 과정에서 실행되는 많은 코드는 native_functions.yaml이라는 스펙으로부터 생성된 코드로, generation_underbar_code.py라는 파일이 이를 처리하여 결과물을 생성합니다. 이 결과물은 Python에서 C++로 넘어가는 과정에서 사용되는 코드의 가장 앞부분에 해당하는 인터페이스를 생성하는 역할을 합니다. 실제 구현체를 보면, THP라는 접두사가 붙어 있는 것을 확인할 수 있는데, 이는 PyTorch 내에서 front-end 함수들을 위한 규약(convention)을 나타냅니다.
native_functions.yaml은 실제 사람이 작성한 코드로, GitHub이나 Python 소스를 빌드하면 여러 가지 방법으로 확인할 수 있는 파일입니다. 이 파일은 aten 라이브러리에 있는 연산자를 정의하고 등록하는 전체 딕셔너리 역할을 하며, dispatcher와 밀접하게 연관되어 다양한 backend나 autograd 관련 기능들을 설정할 수 있습니다. 또한, Python 인터페이스와 C++ 구현을 연결하는 역할도 합니다. 예를 들어, 파일의 맨 위에 있는 add 함수와 그 시그니처 정의, device 체크 여부 등 다양한 특성을 정의하며, 함수가 function인지 method인지도 명시할 수 있습니다. 또한, 어떤 backend들이 dispatch되는지 등의 정보도 이곳에 기술됩니다. 이 정보를 바탕으로 generate_code.py와 같은 도구들이 실제 코드를 생성하게 되며, 이는 PyTorch에서 매우 중요한 부분을 차지합니다.
두 번째로 살펴볼 부분은 실제 matmul operator kernel을 dispatch하는 과정입니다. 이전에 우리는 Python에서 C++로 넘어가는 과정을 다뤘습니다. 이제는 matmul operator에서 kernel을 dispatch하는 과정을 살펴볼 차례입니다. 이 과정에서, 실제로 generic한 일을 하는 front-end 부분이 존재합니다. 이 부분에서는 어떤 함수, 즉 어떤 op를 수행해야 할지를 결정하는데, C++ 수준에서 이 결정을 내리면 그에 맞는 key를 뽑아 dispatcher table에서 해당 backend를 호출하는 과정이 이어집니다.
그 다음으로, 이 과정을 거쳐 실제 matmul operator kernel이 실행됩니다. 여기서 중요한 점은 composite implicit autograd라는 개념입니다. 이 개념은 자동 미분을 위한 기술로, mm operator가 매우 복잡한 작업을 처리하는 데 사용됩니다. 예를 들어, vector 곱하기나 matrix 곱하기와 같은 연산이 있을 수 있습니다. 특히, tensor끼리 곱하기는 여러 matmul operator를 반복해서 수행하는 batched matmul operator로 처리됩니다. 이런 경우, 미분을 별도로 정의하지 않아도 됩니다. 왜냐하면, matmul operator가 결국 미분 가능한 함수들의 시퀀스로 표현될 수 있고, 이를 통해 전체 미분을 처리할 수 있기 때문입니다.
이와 같은 처리 방식을 composite implicit autograd라고 하며, 이 경우 autograd kernel은 자동으로 forward 계산을 수행하는 kernel로 연결됩니다. 해당 코드는 register_composite_implicit_autograd라는 파일에 들어 있으며, 이는 실제로 도구에서 생성된 코드입니다. 이 방식에서는 autograd에 특화된 kernel이 아닌, forward kernel을 수행하고, 그 안에서 필요한 autograd kernel을 자동으로 호출하여 backward 계산이 이루어집니다. 이러한 과정은 인터넷에서 관련 자료를 찾아볼 수 있으며, 오늘은 간단히 설명하고 넘어가겠습니다.
실제로 matmul operator를 처리하는 과정은, aten 아래에 있는 native 디렉터리의 linear_algebra.cpp 파일을 열면 확인할 수 있습니다. 이 파일 안에는 해당 작업을 구현한 코드가 그대로 들어 있습니다. 이 부분은 제가 위 링크에서 복사해 붙여넣은 코드로, 여러 가지 다양한 케이스가 포함되어 있습니다. 특히, mm라고 하는 실제 2D 매트릭스를 계산하는 연산으로 전환되는 부분이 하이라이트된 부분입니다. 이 과정은 linear_algebra.cpp 파일을 열어 보면 구체적으로 확인할 수 있습니다. 말씀드린 것처럼, matmul operator는 다양한 케이스에 대해 일반적으로 동작하도록 설계되어 있으며, 특정 텐서의 차원에 따라 어떤 작업을 수행할지 결정한 후, 이를 처리하기 위해 mm라는 2D 매트릭스 곱셈 연산을 디스패치하게 됩니다.
그럼 dispatch logic이 한 번 더 골라서 결국 mm 작업을 수행하는 커널을 뽑아내게 됩니다. 이 과정은 다시 반복되며, 제가 위에 bold로 하이라이트한 부분이 그 작업을 하는 함수가 호출되는 과정입니다. 여기서 autograd 역할을 하는 커널이 호출되는 것을 볼 수 있습니다. 해당 구현체는 generation 된 코드로, 실제 빌드를 하면 build 디렉토리 안에 있는 VariableType_3 파일에 들어있습니다. 여기서 숫자 3은 코드가 매우 길어 한 파일에 담기 어려워 임의로 나누어져 있기 때문입니다. 이 파일의 세 번째 부분에 있는 함수를 호출하게 됩니다. 이 dispatch 로직은 거의 동일한 방식으로 반복적으로 일어나기 때문에, 이를 반복적인 과정으로 이해할 수 있습니다.
이 부분에서는 autograd kernel에 대해 설명을 드리려고 합니다. 예를 들어, GPU에서 실행할 경우 CUDA 커널을 사용하고, 학습을 위한 함수 호출이 이루어지면 동적 계산 그래프(dynamic computation graph)가 생성됩니다. 이를 통해 CUDA와 autograd 기능이 함께 작동해야 하며, 이 두 기능은 dispatch table에서 적절히 연결됩니다. 여기서 중요한 점은, dispatch key set이 생성되어 우선순위가 가장 높은 키를 선택하여 실행한다는 것입니다. 예를 들어, autograd가 선택되면, autograd kernel이 수행되지만 이 커널은 실제로 backward 계산을 하지 않고, forward 계산을 하면서 동적 계산 그래프를 업데이트하는 역할을 합니다. 이 과정에서 forward 계산은 CUDA kernel로 실제 계산을 수행하는 다른 커널로 리디스패치됩니다.
또한, autograd 커널은 forward 계산을 하면서, backward 계산을 위한 그래프를 자동으로 추가하는 작업을 합니다. 예를 들어, add 연산을 수행할 때 is_input_require_grad 조건을 판별하여, backward 계산을 위한 그래프를 추가하게 됩니다. 이전 구현에서는 native add를 직접 호출했지만, 최신 구현에서는 이를 dispatch 테이블을 통해 재호출하여 수행합니다. 이처럼 autograd kernel은 우선순위가 높은 작업을 처리하며, 하위 작업들을 dispatch하여 수행하도록 설계되어 있습니다. 즉, autograd kernel은 계산을 위한 하위 커널들을 관리하고 실행하는 역할을 합니다.
forward 계산에서 gradient를 직접 계산하지는 않지만, backward를 위한 히스토리를 기록하는 작업이 수행됩니다. 이는 define and run 방식과 define by run 방식의 차이에서 발생하는 문제입니다. define and run 방식은 미리 그래프를 정의하고 이를 기반으로 backward 계산을 하는 방식이지만, PyTorch는 그래프를 명시적으로 정의하지 않고 eager mode에서 계산을 하나하나 수행합니다. 이 과정에서 계산이 완료될 때마다 동적 계산 그래프가 자동으로 생성되고, 후속 backward 계산을 위한 준비가 이루어집니다. eager mode에서는 계산이 즉시 실행되며, 이 계산이 진행되는 동안 필요한 그래프 정보가 계속해서 생성됩니다.
실제로 matmul operator를 처리하는 과정은 aten 아래의 native 디렉터리에 있는 linear_algebra.cpp 파일을 열면 확인할 수 있습니다. 이 파일에는 해당 작업을 구현한 코드가 포함되어 있습니다. 이 부분은 제가 위 링크에서 복사하여 붙여넣은 코드로, 여러 가지 다양한 케이스를 처리하는 내용이 담겨 있습니다. 특히, 2D 매트릭스를 계산하는 연산으로 전환되는 부분이 하이라이트되어 있습니다. 이 과정은 linear_algebra.cpp 파일을 통해 구체적으로 확인할 수 있습니다. 말씀드린 것처럼, matmul operator는 다양한 케이스에 대해 일반적으로 동작하도록 설계되어 있으며, 특정 텐서의 차원을 기반으로 어떤 작업을 수행할지 결정한 후, 이를 처리하기 위해 mm라는 2D 매트릭스 곱셈 연산을 디스패치합니다.
이제 마지막으로 제가 아까 보여드린 call stack의 최상위에서 일어나는 과정에 대해 설명하겠습니다. 여기서 실제로 cublas로 진입하는 과정이 발생하는데, 이 과정에서 'structured'라고 표시된 부분이 있습니다. 이 'structured'는 native_functions.yaml에 있는 'structured'라는 태그와 동일한 의미입니다. 이는 kernel들이 입력과 출력 텐서를 다룰 때, 텐서들을 어떻게 구성하고, 그들의 sanity check을 어떻게 수행할지에 대한 특정 패턴을 정의하고, 그 패턴에 맞춰 구현체가 만들어진다는 것을 의미합니다. 예를 들어, add_out CUDA impl 함수가 호출되기 전에, 이 함수를 래핑하는 다른 함수가 있어 입력 텐서의 sanity를 체크하고, 그 후 structured kernel 구현체가 처리할 수 있는 준비를 마칩니다. 이렇게 준비된 상태에서 개발자는 필요한 부분만 구현하면 됩니다. 그 밑에 있는 껍데기 부분은 자동으로 코드가 생성되고, 잘 정리된 형태의 kernel만 개발자가 구현하면 된다는 생각을 하시면 됩니다. 이 구현체는 아마도 생성된 코드일 것이고, 실제 호출되는 부분은 아마 제대로 된 함수일 것입니다. 이 부분은 'select function'과 관련이 있을 것으로 생각되며, 나중에 확인해 다시 알려드리겠습니다. 오늘은 call stack을 훑어보는 데 대부분의 시간을 소모했는데, 이는 여러분이 알아야 할 거의 대부분의 내용입니다. 나머지 부분은 시간이 허락하면 추가로 설명하겠고, 나머지 부분은 제가 공유한 슬라이드를 참고하여 스스로 공부하시면 좋겠습니다.
지금까지 op에 대해 많은 설명을 드렸습니다. call stack이 실제 op에 대해 찍힌 것이기 때문에 op에 관한 내용은 충분히 다룬 셈입니다. 다만, tensor에 대한 설명은 상대적으로 덜 했으므로, tensor 부분에 대해서 간단히 설명을 덧붙이면 시간이 거의 다 소요될 것 같습니다.
PyTorch의 eager mode 내부에서 tensor는 대부분 operator와 관련된 부분이라고 생각하시면 됩니다. operator에 대해서는 이미 많은 부분을 다뤘으므로, 이제 tensor가 어떻게 구성되는지 살펴보겠습니다. 이 tensor는 멀티 언어 디자인을 고려하여 Python 수준에서 표현되어야 하며, C++에서도 동일한 표현이 필요합니다. 예를 들어, GPU에서 작업을 할 경우, GPU의 device memory에 메모리가 할당되어야 하므로, 이러한 부분들도 고려되어 설계됩니다. 이처럼 tensor는 다양한 요소를 포함한 설계가 필요합니다.
PyTorch의 구조를 설명하면, 위 그림에서 보이는 것처럼 대부분의 부분은 C++로 구성되어 있습니다. 실제로 Python에서 표현되는 tensor는 내부적으로 aten이라는 네임스페이스에 있는 tensor 객체로 들어가게 되며, 이 객체는 주로 impl이라는 멤버를 통해 구현됩니다. impl은 대부분의 정보를 포함하고 있으며, 그 내부에는 storage라는 개념이 존재합니다. storage는 실제 데이터를 저장하는 공간으로, 이는 shared_ptr로 관리되며, 실제 storage impl이라는 구현체가 존재합니다. 예를 들어, GPU에 메모리가 할당되면, 이 메모리 덩어리는 aten 런타임 공간 내의 객체와 1대1로 매핑됩니다. shared_ptr는 특히 view라는 개념이 관련될 때 중요한 역할을 합니다. view를 생성하면 새로운 tensor가 만들어지는 것이 아니라, 기존 storage를 공유하면서 다른 메타데이터를 가진 tensor가 생성됩니다. 이렇게 tensor는 eager mode에서 공유되도록 설계되어 있습니다. 또한, gradient나 관련된 메타데이터, bookkeeping 정보는 tensor 단위에서 저장되도록 되어 있습니다.
Tensor의 라이프사이클은 크게 생성, 소멸, 복사의 세 가지 주요 과정으로 나눌 수 있습니다. Tensor는 두 가지 방법으로 생성될 수 있습니다. 첫 번째는 완전히 새로운 tensor를 처음부터 만드는 경우이고, 두 번째는 기존의 tensor로부터 파생되어 만들어지는 경우입니다. 파생되는 경우에는 기존 tensor와 별개의 storage를 가질 수도 있고, 기존 tensor와 storage를 공유할 수도 있습니다. storage를 공유하는 경우는 제가 아까 언급한 view를 만들 때 발생하며, 이 경우 새로 storage가 생성되는 것이 아니라 shared_ptr로 reference만 공유하게 됩니다. 그 결과 reference count가 증가하게 됩니다. 저는 이러한 방식으로 tensor의 생성 과정을 이해하는 것이 더 명확하다고 생각합니다. 또 다른 방식으로는 tensor의 라이프사이클을 생성 과정으로 해석할 수도 있을 것 같습니다.
소멸 과정은 비교적 간단하지만, 중요한 점은 storage 자체가 소멸되는 것은 아니라는 것입니다. Python과 연결시켜 보면, 대부분의 경우 tensor를 강제로 소멸시키지 않더라도, Python에서의 garbage collection에 의해 자연스럽게 소멸됩니다. Tensor의 Python 구현체가 garbage collection에 의해 소멸되면, 그에 해당하는 C++ 구현체도 함께 소멸되며, 그 안에 포함된 포인터들이 모두 사라집니다. 결국, 최종적으로 storage의 shared_ptr가 소멸되며, 그 shared_ptr가 사라지더라도 하위 구현체는 바로 소멸되지 않고, reference count가 감소하다가 0이 되면, 해당 객체가 참조하고 있던 storage를 소멸시키게 됩니다. 이렇게 tensor의 라이프사이클 소멸이 구현됩니다.
tensor로 할 수 있는 일들은 모두 operator에 해당하며, 이 operator들은 다양한 형태로 나눠질 수 있습니다. 첫 번째로, 입력으로 들어온 객체에 직접 영향을 미치는 in-place operator가 있습니다. 두 번째로, operator 수행 결과로 새로운 tensor가 생성되는 경우가 있는데, 이때 새로운 tensor는 연산 결과로 생성됩니다. 연산 중 값이 크게 변하는 경우에는 새로운 storage가 할당되며, 기존 데이터의 재배치만 일어나는 경우에는 예를 들어 새로운 view가 생기거나 복사가 이루어지기도 합니다. 이러한 경우에는 때로 새로운 storage가 생성되기도 하지만, view의 경우에는 새로운 storage가 생성되지 않습니다. 이렇게 tensor가 할 수 있는 모든 일은 op로 처리되며, 각 op의 동작 방식은 이러한 다양한 특성에 따라 다릅니다. 이러한 구분을 통해 이해하는 것이 저에게는 더 편리했습니다. 이 슬라이드 하단에는 각 경우에 대한 예시를 추가했으니, 궁금하신 분들은 자세히 살펴보시면 좋겠습니다.
view와 storage의 분리는 여러 경우로 나눠질 수 있습니다. 하나는 차원이 변경되는 경우가 있고, 또 다른 경우는 차원이 축소되거나 확장되는 경우입니다. 이와 관련된 다양한 operator들이 존재하며, 이 부분에서 reshape이 포함되지 않은 것을 보면, 아까 말씀하신 내용이 맞다고 생각됩니다. 이 부분은 확인을 해서 코멘트를 달아주시면, 제가 다시 확인 후 답변을 드리겠습니다. 이렇게 해서 storage에 관한 설명은 마쳤고, operator에 대해서는 제가 이미 다룬 개념들이 많아서 오늘은 생략하겠습니다. 궁금하신 분들은 뒷부분을 다시 읽어보시면 좋을 것 같습니다.
오늘 제가 다룬 내용 중에서 커버하지 않은 큰 부분은 바로 runtime에 해당하는 부분입니다. 이 부분은 뒤쪽에 설명이 있지만, 오늘 제가 말씀드린 부분들은 모두 일반적인 내용들이고, 실제로 device-specific한 부분들이 필요합니다. PyTorch는 이러한 device-specific 부분을 처리하기 위해 인터페이스를 정의하고, 각 device별로 해당 클래스를 상속하여 구현하도록 되어 있습니다. 우리가 XPU 같은 새로운 백엔드를 PyTorch에 맞추려고 할 때도, 이와 같은 runtime 구현이 핵심적인 작업이 될 수 있습니다. 이 부분에 대해 궁금한 점이 있으시면, 슬라이드 뒤쪽에 있는 runtime 섹션을 참조하시면 도움이 될 것입니다.