CUDA Graph Capture 기반 학습 루프 - 1차 검증 정리

목적

AICF 에서 CUDA Graph Capture 를 이용한 학습 루프의 최소 단위 검증이 완료되었음을 기록

다음 조건이 구조적으로 성립하는지 확인하는 것

• CUDA Graph capture + replay 환경에서
• forward / backward / optimizer step 이 
• deterministic, 메모리 오염 없이
• 반복 실행 가능함을 검증

 

테스트 범위

• 단순 MLP 구조
• Custom CUDA ops 기반 실생 
  • gemm, relu, relu_bwd, mse_grad, reduce_sum, sgd_step
• Custom autograd 엔진
• Custom SGD optimizer (in-place, capture-safe)
• CUDA Graph
  • cudaStreamBeginCapture
  • cudaGraphInstantiate
  • cudaGraphLaunch 반복 실행

 

테스트 코드 요약

실행 구조

• warmup 단계
  • forward - backward - optimizer
  • 목적
    • 모든 parameter.grad 버퍼 meterialize
    • CUDA allocator 안정화
• CUDA Graph Capture
  • 동일한 학습 step 을 dedicated CUDA stream 에서 capture
  • capture 중
    • 새로운 torch.Tensor 할당 금지
    • torch-side 연산 금지
    • 모든 연산은 AICF op 사용
• Replay 반복
  • cudaGraphLaunch 로 동일 graph 20 회 replay
  • 매 replay 마다
    • parameter update 크기 측정
    • loss_like 관측

 

 

핵심 설계 제약

Capture 중 output tensor 동적 할당 금지

y = backend.op_call("mse_grad", [...])

• 내부에서 torch_empty_like 발생
• 캡처 후 tensor lifetime 보장 불가
• CUDA Graph 가 캡처한 포인터가 replay 시 재사용 / 오염 가능

허용된 방식

backend.op_call_out(
    "mse_grad",
    inputs=[...],
    outputs=[prealloc_dY_buf],  # 캡처 전에 생성
    attrs={}
)

out buffer 실행 모드

 

발견된 문제와 해결

문제 : replay 마다 parameter update 크기가 번갈아 튐

param_step_maxdiff ≈ 6.3e-02
param_step_maxdiff ≈ 1.0e-04
(번갈아 반복)

• SGD step 이 replay 마다 일관되지 않음
• CUDA Graph 내부에서 메모리 alias / lifetime 분괴 발생

원인

• op_call 사용으로 인해
• CUDA Graph capture 중 출력 텐서가 동적으로 생성
• 캡처 종료 후 해당 텐서의 lifetime 이 보장되지 않음
• replay 시 동일 포인터가 다른 버퍼로 재사용됨

해결

• 캡처 중 op_call() 전면 금지
• 모든 capture 경로에서 op_call_out() 만 허용
• mse_grad 출력을 캡처 전에 pre-allocate

 

수정 후 결과

[replay 00] param_step_maxdiff=1.446e-04
[replay 01] param_step_maxdiff=1.437e-04
[replay 02] param_step_maxdiff=1.428e-04
...
[replay 19] param_step_maxdiff=1.283e-04

• parameter update 크기가 연속적이고 안정적으로 감소
• replay 간 step 크기 일관성 확보
• CUDA Graph replay 가 학습 step 1 회로 정상 동작

graph capture 기반 학습 루프의 최소 정확성 검증 완료

 

 

3개의 레이어에서 걸러지는 capture safe 보장 / 로직

 

1) 캡처 중 새 텐서 만들면 안 된다 보장 - aicf_backend.py

여기가 정책을 강제하는 1차 게이트

• capture_begin() / capture_end() 에서 autograd 캡처 가드 토글

from aicf_fw.core.autograd import _set_capture_guard
_set_capture_guard(True/False)

• 가장 중요 : op_call() 에서 캡처 중 호출 금지

from aicf_fw.core.autograd import _IN_CAPTURE_GUARD
if _IN_CAPTURE_GUARD: raise ...

이것이 캡처 중 동적 할당을 구조적으로 막는 핵심, 반대로 op_call _out() 은 출력 버퍼를 외부에서 받기 때문에 캡처 안전

- 캡처 중에는 무조건 out-buffer 경로만 허용한다라는 강제 로직

 

2) leaf grad 포인터가 매 스텝 안정적이다 보장 - autograd.py

학습에서 제일 중요한 안정성 보장 담당

핵심은 leaf grad 에 대해 재바인딩 금지 + 포인터 고정 overwrite 정책이 박혀있다는 점

• 캡처 중 leaf grad 버퍼 새로 만들면 바로 에러
• overwirte 모드에서 leaf grad 는 copy 로 덮어써서 포인터 유지
• accumulate 모드도 add 를 out 형태로 leaf 에 누적

parameter.grad 는 한 번 만들어지면 계속 같은 주소라는 보장

 

3) op 가 어느 CUDA stream 에서 실행되냐 보장 - bindings.cpp

여기가 캡처 / 리플레이의 실행 스트림 결정 담당

• 캡처 중엔 무조건 AICF 전용 stream 사용
• replay 도 같은 aicf_stream 으로 launch

캡처 중 / 리플레이 중 op 런치 스트림이 섞이지 않게 하는 보장

 

 

4) optimizer step 이 pointer-stable 로 수행된다 보장 - optim / sgd...

핵심 in-place update 강제

• inplace=True 면, op_call_out 로 갱신

weight 텐서의 storage 가 재바인딩되지 않는다는 보장

 

5) 실제로 CUDA Graph 가 capture 되녹 replay 되는지 보장 - bindings

• capture_begin - cudaStreamBeginCapture
• capture_end - cudaStreamEndCapture + cudaGraphInstantiate
• replay 에서 cudaGraphLaunch

그래프 캡처 / 인스턴스화 / 런치 파이프라인 자체의 보장

 

• 캡처 중 동적 할당 금지 - aicf_backend(op_call 금지 + guard)
• leaf grad 포인터 안정성 - autograd ( _ensure_leaf_grad_buffer_like, _leaf_overwrite, _leaf_add )
• weight 업데이트 포인터 안정성 - sgd 
• 캡처 / 리플레이 스트림 일관성 - bindings
• 그래프 캡처 자체 - bindings