TorchDynamo

TorchDynamo[link] : Complier frontend

• TorchDynamo is a Python-level JIT (frontend) compiler designed to allow graph compilation
• hooks into the Python frame evaluation API [link] in CPython to dynamically modify Python bytecode right before it is executed
• extract sequences of PyTorch operations into an FX graph (by rewrites Python bytecode)

Process

( AOT autograd $\rightarrow$ recomputation )

AOT autograd

• capturing the backwards pass “ahead-of-time”
• analyzing the forward pass of your model and generating an optimized backward pass ahead of time.

autograd

• autograd[link] : autograd는 데이터(텐서)의 및 실행된 모든 연산들(및 연산 결과가 새로운 텐서인 경우도 포함하여)의 기록을 Function 객체로 구성된 방향성 비순환 그래프(DAG; Directed Acyclic Graph)에 저장(keep)합니다.

Recomputation(aka Activation Checkpointing)

case) cos(cos(a+b+c+d))

saved activation(differentiation) value while forward graph

• naive case saved: a+b+c+d, cos(a+b+c+d), cos(cos(a+b+c+d))
• optimized case return value : a+b+c+d

how choose

• min_cut_rematerialization_partition
• 위 case의 경우, cos는 backward에서 sin(cos)를 계산하는 데 필요하지만, add_2만 넘기고 backward에서 cos를 재연산해도 비용이 저렴.