기존에는 Builder.ops (fwd + bwd + opt 가 섞인 최종 ops) 를 역순회하며 자동 미분 그래프를 생성
목표는
- FWD 를 먼저 최적화 단계에 통과
- 그 결과의 FWD op snapshot 을 기준으로
- emit_bwd 를 호출해 자동 미분 그래프 생성 구조
1. 전체 파이프라인
- FWD build
- FWD optimiz hook ( 현재는 identity)
- FWD ops snapshot
- snapshot 기반 autodiff : emit_bwd 로 BWD graph 생성
- optimizer.step graph 추가
- full compile
- CUDA Graph capture & replay
model.build_backward_after_fwd_opt(loss_vid)- optimize_ir 호출
- fwd_ops_snapshot = list 저장
- build_backward_from_ops 실행
2. Dump 결과 기반 그래프 구조 검증
(A) FWD 빌드 결과
[DUMP] 01_fwd_built__builder.json
num_ops=6
num_values=12
inputs=2
params=4
states=0Op 구성
- gemm x2
- bias_add x2
- relu x1
- cross_entropy x1
(B) BWD 생성 방식 변경 후 결과
[DUMP] 02_bwd_built_post_fwd_opt__builder.json- BWD 가 self.b.ops 저체가 아니라 fwd snapshot 만 역순회하여 생성되었는지
- 실행 결과가 동일하게 수렴 / 정상인지
를 확인
(C) Compile Plan (Full Graph)
[DUMP] 03_compiled_plan__plan.json
num_ops=18
(alias는 이전과 유사하게 적극적으로 설정될 가능성 높음)
kind_id/attr_schema 모두 정상 채워짐최적화 / 컴파일 / 커널 lowering 경로는 정상
3. CUDA Graph 실행 결과
측정 결과
CUDA events (kernel-side)
- Eager: 0.693 ms/step
- CUDA Graph: 0.232 ms/step
Wall time (python + sync 포함)
- Eager: 0.656 ms/step
- CUDA Graph: 0.182 ms/step
해석
- Graph replay는 여전히 Eager 대비 ~3.0–3.6x급 개선
- 이번 실행에서 CUDA event 상 replay가 조금 느려진(0.19→0.23ms) 것은 자연스러운 분산 범위이며,
데이터/클럭/런타임 상태 등 환경 영향으로 흔히 발생한다.
중요한 건:
- 파이프라인 구조 변경(“fwd_opt snapshot → bwd”)이
- CUDA Graph capture/replay 안정성을 깨지 않았고
- 학습 실행이 정상 유지된다는 점이다.