본문 바로가기

AI Compiler framework

AI Framework 전체 실행 구조 문서

AI Framework 는 

  • 모델 구조를 미리 정의하고
  • 실행 중에 실제 계산 경로를 기록하며
  • 그 기록을 거꾸로 재생해 gradient 를 계산하고
  • backend 가 실제 연산을 수행하며
  • optimizer는 계산된 gradient 로 파라미터를 갱신한다

 

1. 전체 시스템의 관찰

User Code
  ↓
Model (Layer 구조)
  ↓
Forward Execution
  ↓
Functional / op_call
  ↓
Backend (Torch / AICF CUDA)
  ↓
Tensor 생성 + Autograd Node 기록
  ↓
loss.backward()
  ↓
Autograd Graph 역순 실행
  ↓
Parameter.grad 채워짐
  ↓
Optimizer.step()

위 흐름이 학습의 전부

 

2. Model / Layer 는 계산을 모른다

Model 의 실제 역할

  • Layer 들을 순서대로 묶어둔 컨테이너
  • forward 호출만 관리
class Model(Module):
    def forward(self, x):
        x = self.l1(x)
        x = self.relu(x)
        x = self.l2(x)
        return x

Model 은 gradient 계산도, backend 도, 커널도 전혀 모른다

 

Layer(Module) 의 역할

  • Parrameter 보관
  • forward 정의
  • functional 호출
class Linear(Module):
    def __init__(self, in_dim, out_dim):
        self.W = Parameter(...)
        self.b = Parameter(...)

    def forward(self, x):
        return functional.linear(x, self.W, self.b)

backward, gradient, 커널 없음

 

3. 진짜 일은 Functional 에서 시작된다.

Functional 이 Framework 의 심장이다

  • backend 에 연산 요청
  • output Tensor 생성
  • autograd Node 생성 및 연결
def linear(x, W, b):
    y = backend.op_call("gemm", [x, W], attrs={...})
    y = backend.op_call("bias_add", [y, b])

    if autograd_enabled:
        node = LinearNode(x, W, b)
        y.creator = node

    return y

forward 계산 + backwrad 준비를 동시에 한다.

 

4. Tensor 의 실체

Tensor =
  data        # 실제 메모리 (CUDA ptr / torch Tensor)
  grad        # 역전파 결과
  creator     # 이 Tensor를 만든 Node

Tensor 는 값이 아니라, 계산 흐름의 노드다.

 

5. Autograd Node 의 정확한 역할

Node 는 연산 하나를 대표한다.

class Node:
    def backward(self, dout):
        raise NotImplementedError

Linear Node

forward 때 저장한 정보 : x, W, b

backward 에서 수행 : dX, dW, dB

class LinearNode(Node):
    def backward(self, dout):
        dx = gemm(dout, self.W.T)
        dW = gemm(self.x.T, dout)
        db = reduce_sum(dout)

        return dx, dW, db

이것이 autograd 의 전부

 

 

6. 계산 그래프는 어떻게 생기나

forward 실행 중 이러한 연결이 만들어짐

x ──▶ LinearNode ──▶ ReLUNode ──▶ LinearNode ──▶ LossNode
  • forward 때 자동으로 생성됨
  • model 구조를 다시 탐색하지 않음
  • 실제 실행된 연산만 기록

if / loop / dynamic shape 가능

 

7. backward 는 되감기 버튼이다

loss.backward 호출 시

  • loss.grad = 1.0 으로 시작
  • creator 를 따라 역순 traversal
  • 각 Node.backward 호출
LossNode.backward
  ↓
LinearNode.backward
  ↓
ReLUNode.backward
  ↓
LinearNode.backward

각 node 는 자기 책임만 수행

 

8. Optimizer 는 단순

Optimizer 는 autograd 를 전혀 모른다

class SGD:
    def step(self):
        for p in model.parameters():
            p.data -= lr * p.grad

그냥 숫자 업데이트

 

9. Backend 의 위치

Backend 는 해당 역할만 수행

  • op_call 실행
  • dtype / variant / kernel 선택
  • 실제 연산 수행
backend.op_call("gemm", inputs, attrs)
  • autograd 로직 없음
  • model / layer 모름
  • CUDA Graph 캡처도 여기서 수행

backend 는 순수 실행 엔진

 

10. Runtime / Graph Capture 는 어디에 붙나

  • python - op_call
  • op_call 내부에서
    • capture 모드면 기록
    • replay 모드면 GraphExec 실행