AICF 통합 학습 그래프 캡처 메커니즘 : AICF 의 Trainstep 완료

1. 테이프 기반의 연산 이력 기록 (Tape Recording)

Model 클래스는 사용자가 정의하는 모든 연산을 _tape 라는 리스트에 기록

• Layer 단위 : Linear, MseLoss 와 같은 복합 연산은 입력, 출력, 그리고 학습될 파라미터 비드를 묶어서 저장한다.
• Op 단위 : add, sub 와 같은 기본 연산은 순수하게 데이터 흐름만 기록

 

2. 역전파 그래프 자동 생성 (AutoDiff Chain)

build_backward(loss-vid) 가 호출되며, 테이프를 역순으로 순회하며 미분 그래프를 그린다.

• 각 레이어는 자신의 emit_backward 를 실행하여 가중치 미분과 입력 미분 연산을 그래프 끝에 덧붙인다.
• 이때, Chain Rule 에 의해 이전 단계에서 계산된 grad_y 가 다음 레이어의 입력으로 전달된다.

 

3. Optimizer 연산 통합 (In-place Update)

미분 그래프 생성이 완료되면, sgd_step 이미터를 사용하여 파라미터 업데이트 연산을 추가

• in_place 설계 : sgd_step 와 같이 입력 파라미터와 출력 파라미터에 동일한 vid 를 할당
• 이는 Planner 가 동일한 메모리 주소를 할당하게 유도하여, 별도의 메모리 복사 없이 GPU 내부에서 즉시 값이 갱신되도록 한다.

 

4. 물리적 메모리 바인딩 (Static Data Binding)

결정적 단계 , CUDA Graph 를 캡처하려면 모든 vid 가 실제 GPU 주소를 가져야 한다.

• Parameter Storage : Model 클래스가 내부 저장소에 실제 Tensor 를 소유한다.
• Full Feed 구성 : 캡처 시점에 사용자 입력 뿐만 아니라 모델이 가진 모든 파라미터 텐서를 sample_feed 로 합쳐서 Executor 에 전달
• KeyError 해결 : 이 과정을 통해 그래프 상의 상징적 변수가 실제 GPU 메모리 포인터와 연결된다.

 

성과 및 의의

• Zero Python Overhead : 학습 루프 내에서 Python 인터프리터의 개입 없이 GPU 가 전 과정을 수행
• Memory Efficiency : 중간 결과값과 그라디언트를 위한 메모리가 캡처 시점에 미리 할당되어 재사용된다.
• Deterministic Training : 동일한 입력에 대해 완전히 동일한 하드웨어 실행 경로를 보장한다.