본문 바로가기

dev_AI_framework

Graph Executor v2 — Forward 및 Training Graph 설계 문서

버전 0.1 — 현재 Forward 그래프 구조 설명과, 학습 전체 단계(Forward + Loss + Backward + Optimizer) 그래프화 확장 설계 요약.


0) 개요

  • 목표: Sequential 구조의 모델을 IR(Intermediate Representation) 형태로 변환 후, Memory Arena 기반의 텐서 메모리 계획 및 관리, CUDA Graph 캡처를 통해 실행 오버헤드를 최소화.
  • 현재 상태: Forward-only. Conv2D, ReLU, BN2D, Flatten, GEMM 등 주요 연산을 IR 노드로 컴파일 후 Arena에서 실행. CuPy 버전별로 다양한 CUDA Graph API 경로를 자동 폴백 처리.
  • 다음 목표: Backward/Optimizer 포함한 학습 전체 단계를 IR 레벨로 확장해, 하나의 CUDA Graph로 전체 학습 스텝을 캡처.

1) 전체 구조

[Model (Sequential)]
  │
  └── GraphCompiler
        ├─ IR (GraphIR)
        │    ├─ TensorDesc: 텐서 shape/dtype/버퍼 위치
        │    └─ Node: 연산 단위 (kind, inputs, outputs, attrs)
        ├─ MemoryArena: 고정 버퍼 할당 및 view 관리
        └─ CUDA Graph Capture: Forward 실행 캡처

2) 주요 구성 요소

(1) TensorDesc

  • 텐서 shape, dtype, buffer_id, offset 정보 보유.
  • 동일 버퍼에서 in-place alias 가능.

(2) Node

  • 그래프 내 연산 단위를 표현.
  • kind(예: "Conv2D"), inputs/outputs 텐서 ID, layer_ref(파라미터 접근용), attrs(Stride, Padding 등) 포함.

(3) MemoryArena

  • GPU 메모리를 1D로 미리 할당하고, TensorDesc를 통해 view로 관리.
  • Forward 및 Backward 모두에서 재할당 없이 사용 가능.

(4) GraphCompiler

  • 모델을 IR로 변환하고 Arena를 구성, CUDA Graph를 캡처.
  • compile()에서 IR 작성 → Arena 생성 → (옵션) 그래프 캡처.
  • run()에서 입력 주입 후 그래프 런치 또는 일반 실행 수행.

3) Forward 그래프 노드 종류

연산 설명

Conv2D Fast-path(groups=1) 시 WS 주입하여 _gconv.forward() 호출
ReLU cp.maximum(X, 0, out=Y) 수행, in-place 가능
BatchNorm2D 간단한 train 모드 BN(mean/var 정규화)
Flatten view 변경만 수행
GEMM(Dense/Linear) gemm_ops.forward() 호출, out 버퍼 직접 기록

4) CUDA Graph 캡처 경로

  1. 고수준 객체 API (Graph.instantiate / upload / launch)
  2. 모듈 함수 API (cp.cuda.graph.upload / launch)
  3. 런타임 포인터 API (cp.cuda.runtime.graphInstantiate / upload / launch)

→ CuPy 버전별 지원 여부에 따라 자동 폴백. stream_ptr을 항상 명시 전달하여 캡처 안전 보장.


5) Forward 그래프의 장점

  • 모든 텐서 shape이 고정되어 런타임 할당 최소화.
  • Arena 기반의 연속 메모리 접근으로 캐시 효율 향상.
  • CUDA Graph 실행으로 Python 오버헤드 제거.
  • ReLU 등 in-place 연산으로 메모리 절약.

6) 학습 그래프 확장 설계

(1) IR 확장

추가 노드 기능

Conv2DBackward dY→dX, dW, db 계산
GEMMBackward dY→dX, dW, db 계산
ActivationBackward mask 기반 ReLU 역전파
BatchNormBackward dγ, dβ 및 dX 계산
SoftmaxCrossEntropy logits→loss, dL/dlogits 계산
Optimizer(SGD/Adam) 파라미터 업데이트 in-place 적용

(2) Arena 확장

  • dX/dW/db/dγ/dβ 텐서까지 포함해 고정 메모리 계획.
  • Forward/Backward용 workspace를 별도로 사전 할당.

(3) 그래프 캡처 범위

Input → Forward → Loss → Backward → Optimizer Step

한 번의 CUDA Graph launch로 전체 학습 스텝 수행.

(4) 캡처 안전성

  • 모든 커널에 현재 stream_ptr 명시(스트림 0 금지)
  • 파라미터 업데이트 포함 모든 연산 GPU in-place
  • host↔device 복사 금지 (입력/출력만 허용)
  • 고정 shape 유지(batch size 변동 시 재캡처 필요)

7) 개발 로드맵

  1. Backward IR 추가 (Conv2D/GEMM/Activation)
  2. Softmax + CE Loss 노드 구현
  3. Optimizer(SGD) 노드 추가
  4. Full Training Graph 캡처 (Forward + Backward + Update)
  5. Gradient Check 및 Eager 모드 비교 검증
  6. Mixed Precision / AMP 확장(fp16 + loss scaling)

8) Forward vs Standalone Training 비교

구분 GraphExecutor (Forward) Mini CNN Train Standalone

실행 구조 IR + Arena + CUDA Graph 파이썬 eager 실행
단계 Forward-only Forward + Backward + Update
메모리 고정 Arena(float32) 동적 cupy.ndarray
성능 높은 Throughput (고정 shape) 유연하지만 상대적으로 느림
확장성 학습 전체 그래프화 가능 검증용 기준으로 유지

9) 결론

현재 GraphExecutor는 Forward-only 단계에서 완성도 높게 동작하며, CUDA Graph를 통한 대폭적인 실행 가속을 제공한다.
다음 단계는 Backward 및 Optimizer를 IR 체계로 통합하여, 학습 전체 단계를 단일 그래프로 캡처하고, 단 한 번의 graphLaunch()로 학습 루프를 수행하는 것이다.